diff --git a/.gitattributes b/.gitattributes
index 0e5fb29a357d4be89184cdca29604b5896650109..be7a13c7ce8640047e0c58bf0f429a200fe51da5 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -143,3 +143,37 @@ dimladder/d512/diagnostics/step_10000/reward_value_curves.png filter=lfs diff=lf
 dimladder/d512/diagnostics/step_15000/reward_value_curves.png filter=lfs diff=lfs merge=lfs -text
 dimladder/d512/diagnostics/step_5000/latent_scatter.png filter=lfs diff=lfs merge=lfs -text
 dimladder/d512/diagnostics/step_5000/reward_value_curves.png filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/10000/default/ocdbt.process_0/d/326c9ee9b3b997eaca5ff01705c1d000 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/10000/default/ocdbt.process_0/d/6ca372f19b981e6b3e3f4d8bfa6918b3 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/10000/default/ocdbt.process_0/d/7bf2ed29052f170443bac5106d28e917 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/10000/default/ocdbt.process_0/d/a81695cfd586dd5c7d6fe409a1e3de3b filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/10000/opt_state/ocdbt.process_0/d/0ac45e219e37878187cab6744dcdde6a filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/10000/opt_state/ocdbt.process_0/d/2a84bced75dfadd31f76b4eff3bbef7c filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/10000/opt_state/ocdbt.process_0/d/cda31d9d73a06b4539bf75d9405cb1a7 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/12500/default/ocdbt.process_0/d/1241673583bfe884858cdcbfa550b66c filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/12500/default/ocdbt.process_0/d/1949e1760c5906a46f5e703e885b8d25 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/12500/default/ocdbt.process_0/d/a575af160125b040bd74df8152f2e347 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/12500/opt_state/ocdbt.process_0/d/088c2713fca76ad7ad42c4b4c8058698 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/12500/opt_state/ocdbt.process_0/d/978ed5aa814340c04f1d9235430548fe filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/12500/opt_state/ocdbt.process_0/d/f617342ff77be531b3a96532d82320dd filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/15000/default/ocdbt.process_0/d/23c6c659c459fc9ce93d7236b9bb1ae9 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/15000/default/ocdbt.process_0/d/30d37c9b93cf1f62022681b916e8f03a filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/15000/default/ocdbt.process_0/d/9b8250888d59e01d229a9438def0f66e filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/15000/default/ocdbt.process_0/d/b342840f343d4de5c616ccf1bfd9bc22 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/15000/opt_state/ocdbt.process_0/d/8f002b592e5173d7e3753a95d485f735 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/15000/opt_state/ocdbt.process_0/d/9b9e27321fed73b1955934f5044312d8 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/15000/opt_state/ocdbt.process_0/d/a867caa8861053a9911611feee61998b filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/15000/opt_state/ocdbt.process_0/d/cc2ff1e06da8bb7e16596a4c5bc212ef filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/5000/default/ocdbt.process_0/d/43bb43789f770ca9e2f6da10d599eaa7 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/5000/default/ocdbt.process_0/d/85c3547ded9ec94ce6e148c4d8551ff5 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/5000/default/ocdbt.process_0/d/9cedb10522084d797d68e95e4dadbbbf filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/5000/default/ocdbt.process_0/d/ff903252b14e2e6439bd63186259473a filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/5000/opt_state/ocdbt.process_0/d/a7d34fd8585833e4964ae3645c104f6a filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/5000/opt_state/ocdbt.process_0/d/b438955b4efbe0aaf8abf4b488aabfc0 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/5000/opt_state/ocdbt.process_0/d/bc5c6aee1e3498353ed7403f9c7eea29 filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/5000/opt_state/ocdbt.process_0/d/f7b7a396546af89363a458e1d243477e filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/diagnostics/step_10000/reward_value_curves.png filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/diagnostics/step_12500/reward_value_curves.png filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/diagnostics/step_15000/reward_value_curves.png filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/diagnostics/step_5000/latent_scatter.png filter=lfs diff=lfs merge=lfs -text
+dimladder/d1024/diagnostics/step_5000/reward_value_curves.png filter=lfs diff=lfs merge=lfs -text
diff --git a/dimladder/d1024/10000/_CHECKPOINT_METADATA b/dimladder/d1024/10000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..b1cc8e6442cf5314fc133cdf4f360056ad919358
--- /dev/null
+++ b/dimladder/d1024/10000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "opt_state": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786620558514317602, "commit_timestamp_nsecs": 1786620559471275162, "custom_metadata": {}}
\ No newline at end of file
diff --git a/dimladder/d1024/10000/default/_METADATA b/dimladder/d1024/10000/default/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..e5b616c5d12ec2aed19a68766f60f4cb366bf6db
--- /dev/null
+++ b/dimladder/d1024/10000/default/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'count')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'key')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('encoder', 'temporal_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('encoder', 'type_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('policy_query',)": {"key_metadata": [{"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('reward_emb', 'pos_emb')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('scaler', 'scaler_params', 'actions', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'actions', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'commands', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'commands', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'proprio_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'proprio_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'rewards', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'rewards', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('value_emb', 'pos_emb')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/dimladder/d1024/10000/default/_sharding b/dimladder/d1024/10000/default/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..5340a724b44e356619cff788181117bffb8958f5
--- /dev/null
+++ b/dimladder/d1024/10000/default/_sharding
@@ -0,0 +1 @@
+{"ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucm5ncy5kZWZhdWx0LmNvdW50":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucm5ncy5kZWZhdWx0LmtleQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci50ZW1wb3JhbF9lbmM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci50eXBlX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5hY3RfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5hY3RfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbm9ybS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbm9ybS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMuYWN0aW9ucy5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMuYWN0aW9ucy5zdGQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucHJvcHJpb19vYnMubWVhbg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucHJvcHJpb19vYnMuc3Rk":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucmV3YXJkcy5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucmV3YXJkcy5zdGQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMudmFsdWVzLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMudmFsdWVzLnN0ZA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X2RlYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X2RlYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X3F1ZXJ5":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2RlYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2RlYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
\ No newline at end of file
diff --git a/dimladder/d1024/10000/default/array_metadatas/process_0 b/dimladder/d1024/10000/default/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..9a849e5ab26571fd7609f10b06d6312c4a32ed8a
--- /dev/null
+++ b/dimladder/d1024/10000/default/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.rngs.default.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.rngs.default.key", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.actions.mean", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.actions.std", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.proprio_obs.mean", "write_shape": [20], "chunk_shape": [20], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.proprio_obs.std", "write_shape": [20], "chunk_shape": [20], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.rewards.mean", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.rewards.std", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.values.mean", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.values.std", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/dimladder/d1024/10000/default/d/17d856f11797651f508f407683651daf b/dimladder/d1024/10000/default/d/17d856f11797651f508f407683651daf
new file mode 100644
index 0000000000000000000000000000000000000000..31729106d2af01ae622a40a19726eebdd64b50eb
Binary files /dev/null and b/dimladder/d1024/10000/default/d/17d856f11797651f508f407683651daf differ
diff --git a/dimladder/d1024/10000/default/manifest.ocdbt b/dimladder/d1024/10000/default/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..d7ec5e654be1bf20970ae883e94c465eb13e084e
Binary files /dev/null and b/dimladder/d1024/10000/default/manifest.ocdbt differ
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/0d07c88e555830a16b82b63a223f12fa b/dimladder/d1024/10000/default/ocdbt.process_0/d/0d07c88e555830a16b82b63a223f12fa
new file mode 100644
index 0000000000000000000000000000000000000000..19803a0275d9ca50d77066188bd45c545a49e505
Binary files /dev/null and b/dimladder/d1024/10000/default/ocdbt.process_0/d/0d07c88e555830a16b82b63a223f12fa differ
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/148143b7b8c098b873031ac7793ef274 b/dimladder/d1024/10000/default/ocdbt.process_0/d/148143b7b8c098b873031ac7793ef274
new file mode 100644
index 0000000000000000000000000000000000000000..9657c84670b083d1c6d448948a884680895a647f
Binary files /dev/null and b/dimladder/d1024/10000/default/ocdbt.process_0/d/148143b7b8c098b873031ac7793ef274 differ
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/1f1b833190c853e1b0f1fd872bb0759e b/dimladder/d1024/10000/default/ocdbt.process_0/d/1f1b833190c853e1b0f1fd872bb0759e
new file mode 100644
index 0000000000000000000000000000000000000000..30e6d3e2ec21151835e12466fffb815b881e2164
Binary files /dev/null and b/dimladder/d1024/10000/default/ocdbt.process_0/d/1f1b833190c853e1b0f1fd872bb0759e differ
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/326c9ee9b3b997eaca5ff01705c1d000 b/dimladder/d1024/10000/default/ocdbt.process_0/d/326c9ee9b3b997eaca5ff01705c1d000
new file mode 100644
index 0000000000000000000000000000000000000000..217365d756f854c385f8c23a1355af6264f5dd5d
--- /dev/null
+++ b/dimladder/d1024/10000/default/ocdbt.process_0/d/326c9ee9b3b997eaca5ff01705c1d000
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c7185fe64856d535cd2bbefc1b5748f8b9f6b2236ee26f4812390175ffc8887
+size 72577024
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/5760caa659b5913497ca0a204f9cebdd b/dimladder/d1024/10000/default/ocdbt.process_0/d/5760caa659b5913497ca0a204f9cebdd
new file mode 100644
index 0000000000000000000000000000000000000000..2c151764fc16f312896b8b3dae10b6bcd31a498e
Binary files /dev/null and b/dimladder/d1024/10000/default/ocdbt.process_0/d/5760caa659b5913497ca0a204f9cebdd differ
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/6ca372f19b981e6b3e3f4d8bfa6918b3 b/dimladder/d1024/10000/default/ocdbt.process_0/d/6ca372f19b981e6b3e3f4d8bfa6918b3
new file mode 100644
index 0000000000000000000000000000000000000000..08d4a9433491ac5acdcf6f43b59a9a68b33654bc
--- /dev/null
+++ b/dimladder/d1024/10000/default/ocdbt.process_0/d/6ca372f19b981e6b3e3f4d8bfa6918b3
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e18757d4288e08faad53867d3b8db4e9808042d1b89bdbda26878662883b83e
+size 77672448
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/7bf2ed29052f170443bac5106d28e917 b/dimladder/d1024/10000/default/ocdbt.process_0/d/7bf2ed29052f170443bac5106d28e917
new file mode 100644
index 0000000000000000000000000000000000000000..df27cf6344f12cc256dd5c302cff362d22be9b38
--- /dev/null
+++ b/dimladder/d1024/10000/default/ocdbt.process_0/d/7bf2ed29052f170443bac5106d28e917
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54daaca40633c9fea7ac75795879ab8deafbad9568669822c7b1210fe0de2ab4
+size 11677696
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/7e12b177a588feaee3e52bc245af086e b/dimladder/d1024/10000/default/ocdbt.process_0/d/7e12b177a588feaee3e52bc245af086e
new file mode 100644
index 0000000000000000000000000000000000000000..3aa055c67cbc70d3324f57c3144d585abc310631
Binary files /dev/null and b/dimladder/d1024/10000/default/ocdbt.process_0/d/7e12b177a588feaee3e52bc245af086e differ
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/a81695cfd586dd5c7d6fe409a1e3de3b b/dimladder/d1024/10000/default/ocdbt.process_0/d/a81695cfd586dd5c7d6fe409a1e3de3b
new file mode 100644
index 0000000000000000000000000000000000000000..988ff267bebaae07297ad082b4df579e15a6c42f
--- /dev/null
+++ b/dimladder/d1024/10000/default/ocdbt.process_0/d/a81695cfd586dd5c7d6fe409a1e3de3b
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d2775497cc4a73658b95e00eac9b8d6587c3ce6036d2725737278a859f3c811
+size 474992640
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/d/c5b41ccde6a64efddc4e39cce373810b b/dimladder/d1024/10000/default/ocdbt.process_0/d/c5b41ccde6a64efddc4e39cce373810b
new file mode 100644
index 0000000000000000000000000000000000000000..abac1b3525f9746316fd336e9a29efb103d6da4c
Binary files /dev/null and b/dimladder/d1024/10000/default/ocdbt.process_0/d/c5b41ccde6a64efddc4e39cce373810b differ
diff --git a/dimladder/d1024/10000/default/ocdbt.process_0/manifest.ocdbt b/dimladder/d1024/10000/default/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..72ba29b5d4991854f6bb99fa373c02e54c7dbc16
Binary files /dev/null and b/dimladder/d1024/10000/default/ocdbt.process_0/manifest.ocdbt differ
diff --git a/dimladder/d1024/10000/opt_state/_METADATA b/dimladder/d1024/10000/opt_state/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..d2e91c3c1ba4762a96dba6939c0102b804fc9dac
--- /dev/null
+++ b/dimladder/d1024/10000/opt_state/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('opt_state', '1', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'temporal_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'type_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('opt_state', '1', '0', 'mu', 'policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('opt_state', '1', '0', 'mu', 'policy_query')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'mu', 'reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'mu', 'value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'temporal_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'type_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('opt_state', '1', '0', 'nu', 'policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('opt_state', '1', '0', 'nu', 'policy_query')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'nu', 'reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'nu', 'value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '1', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('step',)": {"key_metadata": [{"key": "step", "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/dimladder/d1024/10000/opt_state/_sharding b/dimladder/d1024/10000/opt_state/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..17efc14e321f99e13537389be14be30a496b1ce3
--- /dev/null
+++ b/dimladder/d1024/10000/opt_state/_sharding
@@ -0,0 +1 @@
+{"b3B0X3N0YXRlLjEuMC5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9ub3JtLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnR5cGVfZW5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnRlbXBvcmFsX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfcXVlcnk=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9ub3JtLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnR5cGVfZW5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnRlbXBvcmFsX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfcXVlcnk=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMS5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c3RlcA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
\ No newline at end of file
diff --git a/dimladder/d1024/10000/opt_state/array_metadatas/process_0 b/dimladder/d1024/10000/opt_state/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..34e3283d0b868379783226fb713a9a90b9436271
--- /dev/null
+++ b/dimladder/d1024/10000/opt_state/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "opt_state.1.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.1.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/dimladder/d1024/10000/opt_state/d/9a12efa4146996d3545653d6f66a2b6d b/dimladder/d1024/10000/opt_state/d/9a12efa4146996d3545653d6f66a2b6d
new file mode 100644
index 0000000000000000000000000000000000000000..3e9cab3d99ac5a11f240a6819be1e03b2b6651b1
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/d/9a12efa4146996d3545653d6f66a2b6d differ
diff --git a/dimladder/d1024/10000/opt_state/manifest.ocdbt b/dimladder/d1024/10000/opt_state/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..dd0cfb610c5b2fbbb229d513947aec524f4a6db7
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/manifest.ocdbt differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/045fafb8ebe32cd94531a2133a2596f4 b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/045fafb8ebe32cd94531a2133a2596f4
new file mode 100644
index 0000000000000000000000000000000000000000..ad44c59da76f9c7ab56757140897f388bb531439
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/045fafb8ebe32cd94531a2133a2596f4 differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/07a6764747b476d2ec00f3625c332255 b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/07a6764747b476d2ec00f3625c332255
new file mode 100644
index 0000000000000000000000000000000000000000..2d977d2b21220cf2360d39247ad1cb8a7339e95b
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/07a6764747b476d2ec00f3625c332255 differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/0ac45e219e37878187cab6744dcdde6a b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/0ac45e219e37878187cab6744dcdde6a
new file mode 100644
index 0000000000000000000000000000000000000000..bd3c9d49b93a9d1a88becc8ef87b880f17a85995
--- /dev/null
+++ b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/0ac45e219e37878187cab6744dcdde6a
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b9a6515fa618cd482970f8368db8223628139b58de503c8fdbc175870c5c68f
+size 561287168
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/2a84bced75dfadd31f76b4eff3bbef7c b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/2a84bced75dfadd31f76b4eff3bbef7c
new file mode 100644
index 0000000000000000000000000000000000000000..a92920d44b5df28c273a42e9242616067074c0ed
--- /dev/null
+++ b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/2a84bced75dfadd31f76b4eff3bbef7c
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b86e9ef8c27e0bc92302b2e7a76487c29cef460aa0d1dc33b28fe844b5f9466f
+size 59858944
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/3aa40f3a36caf9d86ee80db15c526af5 b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/3aa40f3a36caf9d86ee80db15c526af5
new file mode 100644
index 0000000000000000000000000000000000000000..5600f2fae202333e8f67f49629124d9c5875bcbe
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/3aa40f3a36caf9d86ee80db15c526af5 differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/657a548a295b5c35a5392317fd512763 b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/657a548a295b5c35a5392317fd512763
new file mode 100644
index 0000000000000000000000000000000000000000..40fd8dd639802353fe01786bc46af27ff3318d13
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/657a548a295b5c35a5392317fd512763 differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/7a096585dff6ef326ca97ff277971035 b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/7a096585dff6ef326ca97ff277971035
new file mode 100644
index 0000000000000000000000000000000000000000..b98f5628a3faa25ad0dbf576eaaacc4e4ccb2686
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/7a096585dff6ef326ca97ff277971035 differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/a28af34c14b2e15cb0534cce7d21ad92 b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/a28af34c14b2e15cb0534cce7d21ad92
new file mode 100644
index 0000000000000000000000000000000000000000..e5ab401a84b657af55be19762ac0c1fe4e20fd2a
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/a28af34c14b2e15cb0534cce7d21ad92 differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/c1d30db5a98d30a0f34bbcb3b5333b56 b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/c1d30db5a98d30a0f34bbcb3b5333b56
new file mode 100644
index 0000000000000000000000000000000000000000..68819116d152d7ac5509821d2d00fe74d8195bce
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/c1d30db5a98d30a0f34bbcb3b5333b56 differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/cda31d9d73a06b4539bf75d9405cb1a7 b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/cda31d9d73a06b4539bf75d9405cb1a7
new file mode 100644
index 0000000000000000000000000000000000000000..04239342cbb3d5a1efea376c63ecd87e1e8d1f60
--- /dev/null
+++ b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/cda31d9d73a06b4539bf75d9405cb1a7
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1376b8f122519c9d022d3ccec07436bd092cb0a0179243949f11024a77b0c52
+size 644878336
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/dab8edfdd0919de02840cf87925b523f b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/dab8edfdd0919de02840cf87925b523f
new file mode 100644
index 0000000000000000000000000000000000000000..69f89f27152f9077f2911a2f31e8c39d7eadf9a0
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/dab8edfdd0919de02840cf87925b523f differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/e5feb7bb720c061f5b65091ecdcd5190 b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/e5feb7bb720c061f5b65091ecdcd5190
new file mode 100644
index 0000000000000000000000000000000000000000..41f71228ff8ba61f5ab247abf2f1a5532a0f0ae9
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/d/e5feb7bb720c061f5b65091ecdcd5190 differ
diff --git a/dimladder/d1024/10000/opt_state/ocdbt.process_0/manifest.ocdbt b/dimladder/d1024/10000/opt_state/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..4eae117560aceb72da5dc6cd666a69e36098c529
Binary files /dev/null and b/dimladder/d1024/10000/opt_state/ocdbt.process_0/manifest.ocdbt differ
diff --git a/dimladder/d1024/12500/_CHECKPOINT_METADATA b/dimladder/d1024/12500/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..e4b083f07b3b225412772283d98b01316d582c96
--- /dev/null
+++ b/dimladder/d1024/12500/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "opt_state": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786631820772328638, "commit_timestamp_nsecs": 1786631821800394846, "custom_metadata": {}}
\ No newline at end of file
diff --git a/dimladder/d1024/12500/default/_METADATA b/dimladder/d1024/12500/default/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..e5b616c5d12ec2aed19a68766f60f4cb366bf6db
--- /dev/null
+++ b/dimladder/d1024/12500/default/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'count')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'key')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('encoder', 'temporal_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('encoder', 'type_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('policy_query',)": {"key_metadata": [{"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('reward_emb', 'pos_emb')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('scaler', 'scaler_params', 'actions', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'actions', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'commands', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'commands', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'proprio_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'proprio_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'rewards', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'rewards', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('value_emb', 'pos_emb')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/dimladder/d1024/12500/default/_sharding b/dimladder/d1024/12500/default/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..5340a724b44e356619cff788181117bffb8958f5
--- /dev/null
+++ b/dimladder/d1024/12500/default/_sharding
@@ -0,0 +1 @@
+{"ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucm5ncy5kZWZhdWx0LmNvdW50":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucm5ncy5kZWZhdWx0LmtleQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci50ZW1wb3JhbF9lbmM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci50eXBlX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5hY3RfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5hY3RfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbm9ybS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbm9ybS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMuYWN0aW9ucy5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMuYWN0aW9ucy5zdGQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucHJvcHJpb19vYnMubWVhbg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucHJvcHJpb19vYnMuc3Rk":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucmV3YXJkcy5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucmV3YXJkcy5zdGQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMudmFsdWVzLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMudmFsdWVzLnN0ZA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X2RlYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X2RlYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X3F1ZXJ5":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2RlYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2RlYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
\ No newline at end of file
diff --git a/dimladder/d1024/12500/default/array_metadatas/process_0 b/dimladder/d1024/12500/default/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..9a849e5ab26571fd7609f10b06d6312c4a32ed8a
--- /dev/null
+++ b/dimladder/d1024/12500/default/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.rngs.default.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.rngs.default.key", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.actions.mean", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.actions.std", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.proprio_obs.mean", "write_shape": [20], "chunk_shape": [20], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.proprio_obs.std", "write_shape": [20], "chunk_shape": [20], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.rewards.mean", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.rewards.std", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.values.mean", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.values.std", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/dimladder/d1024/12500/default/d/214e79a4ae85e2637f24a7c584226589 b/dimladder/d1024/12500/default/d/214e79a4ae85e2637f24a7c584226589
new file mode 100644
index 0000000000000000000000000000000000000000..1abe4fb79457467cb5cc4a4a0bc2367507fa751b
Binary files /dev/null and b/dimladder/d1024/12500/default/d/214e79a4ae85e2637f24a7c584226589 differ
diff --git a/dimladder/d1024/12500/default/manifest.ocdbt b/dimladder/d1024/12500/default/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..7c9403145096de6b9017dc79aba75cccc6fbacb3
Binary files /dev/null and b/dimladder/d1024/12500/default/manifest.ocdbt differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/11ba23780060e15ba630dedc2bfbae80 b/dimladder/d1024/12500/default/ocdbt.process_0/d/11ba23780060e15ba630dedc2bfbae80
new file mode 100644
index 0000000000000000000000000000000000000000..c9d739c6f8fff2f8bdf9eef25fb9ec13cefc401e
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/11ba23780060e15ba630dedc2bfbae80 differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/1241673583bfe884858cdcbfa550b66c b/dimladder/d1024/12500/default/ocdbt.process_0/d/1241673583bfe884858cdcbfa550b66c
new file mode 100644
index 0000000000000000000000000000000000000000..12630baa984abb1c3eb6bf540447046dea4ca55a
--- /dev/null
+++ b/dimladder/d1024/12500/default/ocdbt.process_0/d/1241673583bfe884858cdcbfa550b66c
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:057f032e1d864d0bc12c9793775399d6c1849fd9de29a613a0c733a2cb768c85
+size 31162368
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/1949e1760c5906a46f5e703e885b8d25 b/dimladder/d1024/12500/default/ocdbt.process_0/d/1949e1760c5906a46f5e703e885b8d25
new file mode 100644
index 0000000000000000000000000000000000000000..5e9d8b203fc84bc0d74405e158b752376b4524f2
--- /dev/null
+++ b/dimladder/d1024/12500/default/ocdbt.process_0/d/1949e1760c5906a46f5e703e885b8d25
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc5d924a9566187fc799c730d8224ac642fc0a741a54279b0b97d7ff88c978f4
+size 392151040
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/28d40fcfdd8a3de5b68638fae922d3d7 b/dimladder/d1024/12500/default/ocdbt.process_0/d/28d40fcfdd8a3de5b68638fae922d3d7
new file mode 100644
index 0000000000000000000000000000000000000000..de83f0b9b10ef1890ff64f7cfb4359402ceeb526
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/28d40fcfdd8a3de5b68638fae922d3d7 differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/29258bb399c512174d68f6c48abdb36d b/dimladder/d1024/12500/default/ocdbt.process_0/d/29258bb399c512174d68f6c48abdb36d
new file mode 100644
index 0000000000000000000000000000000000000000..2a125bfe0a5e91b2adfd386ab884b684e2606751
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/29258bb399c512174d68f6c48abdb36d differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/751b60e00c0a8cc7ec6874178e135cef b/dimladder/d1024/12500/default/ocdbt.process_0/d/751b60e00c0a8cc7ec6874178e135cef
new file mode 100644
index 0000000000000000000000000000000000000000..c215f90ab7a9104f65aa5f51603a3c7c1562b264
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/751b60e00c0a8cc7ec6874178e135cef differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/91608df9dbafe6107af52439b783049f b/dimladder/d1024/12500/default/ocdbt.process_0/d/91608df9dbafe6107af52439b783049f
new file mode 100644
index 0000000000000000000000000000000000000000..1dd0d34577a1fe9cdf70361250c7a8e20c053028
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/91608df9dbafe6107af52439b783049f differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/a575af160125b040bd74df8152f2e347 b/dimladder/d1024/12500/default/ocdbt.process_0/d/a575af160125b040bd74df8152f2e347
new file mode 100644
index 0000000000000000000000000000000000000000..09ead4ca6c1bef5f1768a5df222c0b3c80d8bc9c
--- /dev/null
+++ b/dimladder/d1024/12500/default/ocdbt.process_0/d/a575af160125b040bd74df8152f2e347
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23458424210c5a44aa997581a0f261cce56535850b0a46c1ae42e89c2c62cc60
+size 213671936
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/c8960ab8a967ec802b31d137195daa0d b/dimladder/d1024/12500/default/ocdbt.process_0/d/c8960ab8a967ec802b31d137195daa0d
new file mode 100644
index 0000000000000000000000000000000000000000..d42b521280e215fa72b1011933c76815ab90c692
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/c8960ab8a967ec802b31d137195daa0d differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/da32ae6054f749f3bdad6d94f936280d b/dimladder/d1024/12500/default/ocdbt.process_0/d/da32ae6054f749f3bdad6d94f936280d
new file mode 100644
index 0000000000000000000000000000000000000000..756dce5d7dd3d9bb2bc2f84e12ac265d41dbb308
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/da32ae6054f749f3bdad6d94f936280d differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/e64ca582adb54c0cd029a6495598c879 b/dimladder/d1024/12500/default/ocdbt.process_0/d/e64ca582adb54c0cd029a6495598c879
new file mode 100644
index 0000000000000000000000000000000000000000..2c151764fc16f312896b8b3dae10b6bcd31a498e
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/e64ca582adb54c0cd029a6495598c879 differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/e6775bab51dcbdec52025f1ad60b4b3f b/dimladder/d1024/12500/default/ocdbt.process_0/d/e6775bab51dcbdec52025f1ad60b4b3f
new file mode 100644
index 0000000000000000000000000000000000000000..9657c84670b083d1c6d448948a884680895a647f
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/e6775bab51dcbdec52025f1ad60b4b3f differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/d/fbcc172b5e32992348c0da3ecd3e5ad0 b/dimladder/d1024/12500/default/ocdbt.process_0/d/fbcc172b5e32992348c0da3ecd3e5ad0
new file mode 100644
index 0000000000000000000000000000000000000000..2e67bbe7fa19fc6e1297e3deb06e9e2568ab4eb3
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/d/fbcc172b5e32992348c0da3ecd3e5ad0 differ
diff --git a/dimladder/d1024/12500/default/ocdbt.process_0/manifest.ocdbt b/dimladder/d1024/12500/default/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..c5e3031ee0b35df3d5b06f0b157188d05a468e8f
Binary files /dev/null and b/dimladder/d1024/12500/default/ocdbt.process_0/manifest.ocdbt differ
diff --git a/dimladder/d1024/12500/opt_state/_METADATA b/dimladder/d1024/12500/opt_state/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..d2e91c3c1ba4762a96dba6939c0102b804fc9dac
--- /dev/null
+++ b/dimladder/d1024/12500/opt_state/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('opt_state', '1', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'temporal_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'type_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('opt_state', '1', '0', 'mu', 'policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('opt_state', '1', '0', 'mu', 'policy_query')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'mu', 'reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'mu', 'value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'temporal_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'type_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('opt_state', '1', '0', 'nu', 'policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('opt_state', '1', '0', 'nu', 'policy_query')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'nu', 'reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'nu', 'value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '1', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('step',)": {"key_metadata": [{"key": "step", "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/dimladder/d1024/12500/opt_state/_sharding b/dimladder/d1024/12500/opt_state/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..17efc14e321f99e13537389be14be30a496b1ce3
--- /dev/null
+++ b/dimladder/d1024/12500/opt_state/_sharding
@@ -0,0 +1 @@
+{"b3B0X3N0YXRlLjEuMC5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9ub3JtLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnR5cGVfZW5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnRlbXBvcmFsX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfcXVlcnk=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9ub3JtLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnR5cGVfZW5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnRlbXBvcmFsX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfcXVlcnk=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMS5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c3RlcA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
\ No newline at end of file
diff --git a/dimladder/d1024/12500/opt_state/array_metadatas/process_0 b/dimladder/d1024/12500/opt_state/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..34e3283d0b868379783226fb713a9a90b9436271
--- /dev/null
+++ b/dimladder/d1024/12500/opt_state/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "opt_state.1.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.1.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/dimladder/d1024/12500/opt_state/d/707e8e1714565b9cfd5a3500ab108341 b/dimladder/d1024/12500/opt_state/d/707e8e1714565b9cfd5a3500ab108341
new file mode 100644
index 0000000000000000000000000000000000000000..39075df52a9e293dcfc3ba960ff430617fa8a12d
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/d/707e8e1714565b9cfd5a3500ab108341 differ
diff --git a/dimladder/d1024/12500/opt_state/manifest.ocdbt b/dimladder/d1024/12500/opt_state/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..f6096cce2761979727e005f0f55b750672319c60
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/manifest.ocdbt differ
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/00b46f6cd5198c6c7d17deb57f54fb8c b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/00b46f6cd5198c6c7d17deb57f54fb8c
new file mode 100644
index 0000000000000000000000000000000000000000..b6c0b15d47bd19d9879460484060ff0b73aa601d
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/00b46f6cd5198c6c7d17deb57f54fb8c differ
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/088c2713fca76ad7ad42c4b4c8058698 b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/088c2713fca76ad7ad42c4b4c8058698
new file mode 100644
index 0000000000000000000000000000000000000000..89939c7866fa3ba3140fb2228e1a81d952bea3a4
--- /dev/null
+++ b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/088c2713fca76ad7ad42c4b4c8058698
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4108599361388b67829b50fb56b372906b111d17328ae9ff727ca06a99f356d8
+size 420360192
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/419cdd89caa379471737b4b45905c17a b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/419cdd89caa379471737b4b45905c17a
new file mode 100644
index 0000000000000000000000000000000000000000..b561264d90f01610e37aaa2488633a3f0bf54a01
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/419cdd89caa379471737b4b45905c17a differ
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/942ea2b44b27254a7ddae11ddd163919 b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/942ea2b44b27254a7ddae11ddd163919
new file mode 100644
index 0000000000000000000000000000000000000000..9910617f0a8a1cf4e352b6552e0ef161029b30d9
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/942ea2b44b27254a7ddae11ddd163919 differ
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/978ed5aa814340c04f1d9235430548fe b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/978ed5aa814340c04f1d9235430548fe
new file mode 100644
index 0000000000000000000000000000000000000000..88e78ef27e52b526445c96ae99386675612a35ad
--- /dev/null
+++ b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/978ed5aa814340c04f1d9235430548fe
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5aa3c3b3db1d347edae1f6dce40559a538113ed13512708d460bd1ec4727331
+size 11923456
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/a62407cb983c6f79ba404310f3beebfd b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/a62407cb983c6f79ba404310f3beebfd
new file mode 100644
index 0000000000000000000000000000000000000000..41f71228ff8ba61f5ab247abf2f1a5532a0f0ae9
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/a62407cb983c6f79ba404310f3beebfd differ
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/aabe09c43fd275a600d97b3828c77165 b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/aabe09c43fd275a600d97b3828c77165
new file mode 100644
index 0000000000000000000000000000000000000000..d621163d62d6fb6bc271f39d902f0ded7c12754e
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/aabe09c43fd275a600d97b3828c77165 differ
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/bd3947a72ed97f4d5f090a7519d0d577 b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/bd3947a72ed97f4d5f090a7519d0d577
new file mode 100644
index 0000000000000000000000000000000000000000..57525fb53f70657530ada679f575359a8bafc450
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/bd3947a72ed97f4d5f090a7519d0d577 differ
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/c97beb5c3f131460d391a85b12279fb0 b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/c97beb5c3f131460d391a85b12279fb0
new file mode 100644
index 0000000000000000000000000000000000000000..68819116d152d7ac5509821d2d00fe74d8195bce
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/c97beb5c3f131460d391a85b12279fb0 differ
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/f617342ff77be531b3a96532d82320dd b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/f617342ff77be531b3a96532d82320dd
new file mode 100644
index 0000000000000000000000000000000000000000..ab640e6f75409c70cea85f4fd2f4e5c2485342b3
--- /dev/null
+++ b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/f617342ff77be531b3a96532d82320dd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a793fc19f19fb0137544f7f749abaa7f5a598dba4d41ecaf6165b5578306058c
+size 838623232
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/fff7eb0e37bfecab290aba9fa038d45f b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/fff7eb0e37bfecab290aba9fa038d45f
new file mode 100644
index 0000000000000000000000000000000000000000..94003efa6fcda0c32f723e6ad855802833ee9cdc
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/ocdbt.process_0/d/fff7eb0e37bfecab290aba9fa038d45f differ
diff --git a/dimladder/d1024/12500/opt_state/ocdbt.process_0/manifest.ocdbt b/dimladder/d1024/12500/opt_state/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..1119dd607cd20b5352cb5e7d23f1701842503a97
Binary files /dev/null and b/dimladder/d1024/12500/opt_state/ocdbt.process_0/manifest.ocdbt differ
diff --git a/dimladder/d1024/15000/_CHECKPOINT_METADATA b/dimladder/d1024/15000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..1298260f6410b73793ae6b49a9d521587e9af608
--- /dev/null
+++ b/dimladder/d1024/15000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "opt_state": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786633014231728586, "commit_timestamp_nsecs": 1786633015376668698, "custom_metadata": {}}
\ No newline at end of file
diff --git a/dimladder/d1024/15000/default/_METADATA b/dimladder/d1024/15000/default/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..e5b616c5d12ec2aed19a68766f60f4cb366bf6db
--- /dev/null
+++ b/dimladder/d1024/15000/default/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'count')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'key')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('encoder', 'temporal_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('encoder', 'type_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('policy_query',)": {"key_metadata": [{"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('reward_emb', 'pos_emb')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('scaler', 'scaler_params', 'actions', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'actions', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'commands', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'commands', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'proprio_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'proprio_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'rewards', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'rewards', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('value_emb', 'pos_emb')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/dimladder/d1024/15000/default/_sharding b/dimladder/d1024/15000/default/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..5340a724b44e356619cff788181117bffb8958f5
--- /dev/null
+++ b/dimladder/d1024/15000/default/_sharding
@@ -0,0 +1 @@
+{"ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucm5ncy5kZWZhdWx0LmNvdW50":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucm5ncy5kZWZhdWx0LmtleQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci50ZW1wb3JhbF9lbmM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci50eXBlX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5hY3RfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5hY3RfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbm9ybS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbm9ybS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMuYWN0aW9ucy5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMuYWN0aW9ucy5zdGQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucHJvcHJpb19vYnMubWVhbg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucHJvcHJpb19vYnMuc3Rk":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucmV3YXJkcy5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucmV3YXJkcy5zdGQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMudmFsdWVzLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMudmFsdWVzLnN0ZA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X2RlYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X2RlYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X3F1ZXJ5":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2RlYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2RlYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
\ No newline at end of file
diff --git a/dimladder/d1024/15000/default/array_metadatas/process_0 b/dimladder/d1024/15000/default/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..a029df478826e9cd7e6018a0676d0048bb206b95
--- /dev/null
+++ b/dimladder/d1024/15000/default/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.rngs.default.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.rngs.default.key", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.actions.mean", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.actions.std", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.proprio_obs.mean", "write_shape": [20], "chunk_shape": [20], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.proprio_obs.std", "write_shape": [20], "chunk_shape": [20], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.rewards.mean", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.rewards.std", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.values.mean", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.values.std", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/dimladder/d1024/15000/default/d/262e105b032c21c56dbf5c39fd882924 b/dimladder/d1024/15000/default/d/262e105b032c21c56dbf5c39fd882924
new file mode 100644
index 0000000000000000000000000000000000000000..a075719d114121a08a02f5757e9b1b8497b5a641
Binary files /dev/null and b/dimladder/d1024/15000/default/d/262e105b032c21c56dbf5c39fd882924 differ
diff --git a/dimladder/d1024/15000/default/manifest.ocdbt b/dimladder/d1024/15000/default/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..c23dde637b0996fb48c544fd6138c86e1e9ac28b
Binary files /dev/null and b/dimladder/d1024/15000/default/manifest.ocdbt differ
diff --git a/dimladder/d1024/15000/default/ocdbt.process_0/d/23c6c659c459fc9ce93d7236b9bb1ae9 b/dimladder/d1024/15000/default/ocdbt.process_0/d/23c6c659c459fc9ce93d7236b9bb1ae9
new file mode 100644
index 0000000000000000000000000000000000000000..f67b6eb1c2177d0eda6827fed988169812078af0
--- /dev/null
+++ b/dimladder/d1024/15000/default/ocdbt.process_0/d/23c6c659c459fc9ce93d7236b9bb1ae9
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d53c8627c6b352028ecc10350012daaadb52111d9fad45d3e1684e27bf7000f7
+size 86544384
diff --git a/dimladder/d1024/15000/default/ocdbt.process_0/d/2e0fc2e896e3b47a19b9816283e65e5b b/dimladder/d1024/15000/default/ocdbt.process_0/d/2e0fc2e896e3b47a19b9816283e65e5b
new file mode 100644
index 0000000000000000000000000000000000000000..50fc40ef8ccc35a1fc22b71698c1dee6281d38a5
Binary files /dev/null and b/dimladder/d1024/15000/default/ocdbt.process_0/d/2e0fc2e896e3b47a19b9816283e65e5b differ
diff --git a/dimladder/d1024/15000/default/ocdbt.process_0/d/30d37c9b93cf1f62022681b916e8f03a b/dimladder/d1024/15000/default/ocdbt.process_0/d/30d37c9b93cf1f62022681b916e8f03a
new file mode 100644
index 0000000000000000000000000000000000000000..28a1050808d6bd91389aba15c299f7b1cd7c44d3
--- /dev/null
+++ b/dimladder/d1024/15000/default/ocdbt.process_0/d/30d37c9b93cf1f62022681b916e8f03a
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c5e944640a0de84ae2fd56a3b26ac677172b7ef32fc4af180a8aaf1e3514ebb
+size 492187648
diff --git a/dimladder/d1024/15000/default/ocdbt.process_0/d/9b8250888d59e01d229a9438def0f66e b/dimladder/d1024/15000/default/ocdbt.process_0/d/9b8250888d59e01d229a9438def0f66e
new file mode 100644
index 0000000000000000000000000000000000000000..a539d638acf284722300cbd9877d198fa98f3b0c
--- /dev/null
+++ b/dimladder/d1024/15000/default/ocdbt.process_0/d/9b8250888d59e01d229a9438def0f66e
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ab58c14d05317a7426b645617300263de8ce10a8cfd8e813323c078f33609bf
+size 57339904
diff --git a/dimladder/d1024/15000/default/ocdbt.process_0/d/9e69ec396a9ec95eeb745b65dcc0a2c6 b/dimladder/d1024/15000/default/ocdbt.process_0/d/9e69ec396a9ec95eeb745b65dcc0a2c6
new file mode 100644
index 0000000000000000000000000000000000000000..9657c84670b083d1c6d448948a884680895a647f
Binary files /dev/null and b/dimladder/d1024/15000/default/ocdbt.process_0/d/9e69ec396a9ec95eeb745b65dcc0a2c6 differ
diff --git a/dimladder/d1024/15000/default/ocdbt.process_0/d/b342840f343d4de5c616ccf1bfd9bc22 b/dimladder/d1024/15000/default/ocdbt.process_0/d/b342840f343d4de5c616ccf1bfd9bc22
new file mode 100644
index 0000000000000000000000000000000000000000..7312b0a8bee631b6623771f96efb1eecd7047583
--- /dev/null
+++ b/dimladder/d1024/15000/default/ocdbt.process_0/d/b342840f343d4de5c616ccf1bfd9bc22
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:986d4af345189107ea3945decc6feb0098ad62e665159caa7720b7942345fb17
+size 1015439
diff --git a/dimladder/d1024/15000/default/ocdbt.process_0/d/de999ea9475300f05a608b3e17a76f1b b/dimladder/d1024/15000/default/ocdbt.process_0/d/de999ea9475300f05a608b3e17a76f1b
new file mode 100644
index 0000000000000000000000000000000000000000..e1433eb9bd53bf231913c66c9fa08d04d2397b33
Binary files /dev/null and b/dimladder/d1024/15000/default/ocdbt.process_0/d/de999ea9475300f05a608b3e17a76f1b differ
diff --git a/dimladder/d1024/15000/default/ocdbt.process_0/manifest.ocdbt b/dimladder/d1024/15000/default/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..c179c395af1540050d4e20efe38c7fa98c9426a0
Binary files /dev/null and b/dimladder/d1024/15000/default/ocdbt.process_0/manifest.ocdbt differ
diff --git a/dimladder/d1024/15000/opt_state/_METADATA b/dimladder/d1024/15000/opt_state/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..d2e91c3c1ba4762a96dba6939c0102b804fc9dac
--- /dev/null
+++ b/dimladder/d1024/15000/opt_state/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('opt_state', '1', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'temporal_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'type_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('opt_state', '1', '0', 'mu', 'policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('opt_state', '1', '0', 'mu', 'policy_query')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'mu', 'reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'mu', 'value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'temporal_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'type_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('opt_state', '1', '0', 'nu', 'policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('opt_state', '1', '0', 'nu', 'policy_query')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'nu', 'reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'nu', 'value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '1', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('step',)": {"key_metadata": [{"key": "step", "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/dimladder/d1024/15000/opt_state/_sharding b/dimladder/d1024/15000/opt_state/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..17efc14e321f99e13537389be14be30a496b1ce3
--- /dev/null
+++ b/dimladder/d1024/15000/opt_state/_sharding
@@ -0,0 +1 @@
+{"b3B0X3N0YXRlLjEuMC5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9ub3JtLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnR5cGVfZW5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnRlbXBvcmFsX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfcXVlcnk=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9ub3JtLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnR5cGVfZW5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnRlbXBvcmFsX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfcXVlcnk=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMS5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c3RlcA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
\ No newline at end of file
diff --git a/dimladder/d1024/15000/opt_state/array_metadatas/process_0 b/dimladder/d1024/15000/opt_state/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..34e3283d0b868379783226fb713a9a90b9436271
--- /dev/null
+++ b/dimladder/d1024/15000/opt_state/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "opt_state.1.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.1.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/dimladder/d1024/15000/opt_state/d/7d6c76beb490ee4e8c0a8ecaec071108 b/dimladder/d1024/15000/opt_state/d/7d6c76beb490ee4e8c0a8ecaec071108
new file mode 100644
index 0000000000000000000000000000000000000000..649a912feb094d6199744f1b851c49b6af79c773
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/d/7d6c76beb490ee4e8c0a8ecaec071108 differ
diff --git a/dimladder/d1024/15000/opt_state/manifest.ocdbt b/dimladder/d1024/15000/opt_state/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..693243b483358f4a95885962625303b3b1c800fd
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/manifest.ocdbt differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/11fe639ab317f2f18c919c7d5343527c b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/11fe639ab317f2f18c919c7d5343527c
new file mode 100644
index 0000000000000000000000000000000000000000..3f4a96b34555eae59fbb63362845201b36a8e643
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/11fe639ab317f2f18c919c7d5343527c differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/2302a3151bf2ffc9c4ca2e5bed539c96 b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/2302a3151bf2ffc9c4ca2e5bed539c96
new file mode 100644
index 0000000000000000000000000000000000000000..54f49bdd3cdd5e2ab3cd57c5b3f9a7763620611a
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/2302a3151bf2ffc9c4ca2e5bed539c96 differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/2fed33b77f2ae910db93ac0904d06c75 b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/2fed33b77f2ae910db93ac0904d06c75
new file mode 100644
index 0000000000000000000000000000000000000000..3a142225b5c34cc51883a6b4d48d0f95c47889d3
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/2fed33b77f2ae910db93ac0904d06c75 differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/5b2f1c389906d78f75e80c9e6ca2c599 b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/5b2f1c389906d78f75e80c9e6ca2c599
new file mode 100644
index 0000000000000000000000000000000000000000..41f71228ff8ba61f5ab247abf2f1a5532a0f0ae9
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/5b2f1c389906d78f75e80c9e6ca2c599 differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/82c8a03b85709dc591ad5fbe02cdf986 b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/82c8a03b85709dc591ad5fbe02cdf986
new file mode 100644
index 0000000000000000000000000000000000000000..68819116d152d7ac5509821d2d00fe74d8195bce
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/82c8a03b85709dc591ad5fbe02cdf986 differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/8f002b592e5173d7e3753a95d485f735 b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/8f002b592e5173d7e3753a95d485f735
new file mode 100644
index 0000000000000000000000000000000000000000..ffa94abdfdd98ced5218cf3abc8567922bdfe613
--- /dev/null
+++ b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/8f002b592e5173d7e3753a95d485f735
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf5d9cf5fabe750f9ac1e63626970cde58cf84a357509d02691312e08db29da1
+size 1118453760
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/97635d7d25b94faafac94c1436c3307d b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/97635d7d25b94faafac94c1436c3307d
new file mode 100644
index 0000000000000000000000000000000000000000..8e2a2a696b6c36c281966bbe3cd74bc5f85a8a32
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/97635d7d25b94faafac94c1436c3307d differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/9b9e27321fed73b1955934f5044312d8 b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/9b9e27321fed73b1955934f5044312d8
new file mode 100644
index 0000000000000000000000000000000000000000..a1c1a196cf6cda6248a9dbb01448716d10a561dd
--- /dev/null
+++ b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/9b9e27321fed73b1955934f5044312d8
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a13ad97b63184b65c87ffa2e57e5a7e37510ca6f6d313080814672792846704
+size 92217344
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/a867caa8861053a9911611feee61998b b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/a867caa8861053a9911611feee61998b
new file mode 100644
index 0000000000000000000000000000000000000000..8e639a58094fbbe785ad5e9ba0048d13b98cd17a
--- /dev/null
+++ b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/a867caa8861053a9911611feee61998b
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:442cad043236eabb73c6e4ac29ca9ff34cd9c94f0a793113eeb8b4f3f0b6632a
+size 109636
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/b023b70debb27930713b5104d32a76f9 b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/b023b70debb27930713b5104d32a76f9
new file mode 100644
index 0000000000000000000000000000000000000000..a81f257b4a4f7122ae20a78005cc2244c0a0aa8e
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/b023b70debb27930713b5104d32a76f9 differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/cc2ff1e06da8bb7e16596a4c5bc212ef b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/cc2ff1e06da8bb7e16596a4c5bc212ef
new file mode 100644
index 0000000000000000000000000000000000000000..570de7c25135d780573c384daef3ee53c252012a
--- /dev/null
+++ b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/cc2ff1e06da8bb7e16596a4c5bc212ef
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:012167f02834247b9e0c29788bc138d04ccb8b6a55486a90e0b2e8fd4ee99b35
+size 61669376
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/ec949c755c225c6f5f23888493189157 b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/ec949c755c225c6f5f23888493189157
new file mode 100644
index 0000000000000000000000000000000000000000..ff2d3b9b3f90e3421325c825844119a16f0490fe
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/ec949c755c225c6f5f23888493189157 differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/fac8a669095aac8d7cbe59ce5ee5d463 b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/fac8a669095aac8d7cbe59ce5ee5d463
new file mode 100644
index 0000000000000000000000000000000000000000..e5ab401a84b657af55be19762ac0c1fe4e20fd2a
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/d/fac8a669095aac8d7cbe59ce5ee5d463 differ
diff --git a/dimladder/d1024/15000/opt_state/ocdbt.process_0/manifest.ocdbt b/dimladder/d1024/15000/opt_state/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..b695cf826943d84005996708e7f62e6484fa36af
Binary files /dev/null and b/dimladder/d1024/15000/opt_state/ocdbt.process_0/manifest.ocdbt differ
diff --git a/dimladder/d1024/5000/_CHECKPOINT_METADATA b/dimladder/d1024/5000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..58907d10391df5e98cff6825b774c0f5ab746545
--- /dev/null
+++ b/dimladder/d1024/5000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "opt_state": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786617925271087042, "commit_timestamp_nsecs": 1786617926605721417, "custom_metadata": {}}
\ No newline at end of file
diff --git a/dimladder/d1024/5000/default/_METADATA b/dimladder/d1024/5000/default/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..e5b616c5d12ec2aed19a68766f60f4cb366bf6db
--- /dev/null
+++ b/dimladder/d1024/5000/default/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'count')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('dynamics', 'layers', '0', 'attention_1', 'rngs', 'default', 'key')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "rngs", "key_type": 2}, {"key": "default", "key_type": 2}, {"key": "key", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'mean')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'var')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('encoder', 'temporal_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('encoder', 'type_enc')": {"key_metadata": [{"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('policy_query',)": {"key_metadata": [{"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('reward_emb', 'pos_emb')": {"key_metadata": [{"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('scaler', 'scaler_params', 'actions', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'actions', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "actions", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('scaler', 'scaler_params', 'commands', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'commands', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "commands", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'extero_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "extero_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('scaler', 'scaler_params', 'proprio_obs', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'proprio_obs', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "proprio_obs", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20]}}, "('scaler', 'scaler_params', 'rewards', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'rewards', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "rewards", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'mean')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "mean", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('scaler', 'scaler_params', 'values', 'std')": {"key_metadata": [{"key": "scaler", "key_type": 2}, {"key": "scaler_params", "key_type": 2}, {"key": "values", "key_type": 2}, {"key": "std", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('value_emb', 'pos_emb')": {"key_metadata": [{"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/dimladder/d1024/5000/default/_sharding b/dimladder/d1024/5000/default/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..5340a724b44e356619cff788181117bffb8958f5
--- /dev/null
+++ b/dimladder/d1024/5000/default/_sharding
@@ -0,0 +1 @@
+{"ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucm5ncy5kZWZhdWx0LmNvdW50":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEucm5ncy5kZWZhdWx0LmtleQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZHluYW1pY3MubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci50ZW1wb3JhbF9lbmM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci50eXBlX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5hY3RfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5hY3RfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXI0LjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIxLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIyLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi52YXI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9ibi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjAuZG93bnNhbXBsZV9jb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuY29udjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuY29udjIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4xLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5pbWFnZV9lbmMudHJ1bmtzLjAubGF5ZXIzLjEuYm4yLnZhcg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbWxwLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbm9ybS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5sYXRlbnRfbm9ybS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZW5jb2Rlci5wcm9wcmlvX29ic19wcm9qLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","ZnV0X2FjdHNfZW1iZWQucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMuYWN0aW9ucy5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMuYWN0aW9ucy5zdGQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucHJvcHJpb19vYnMubWVhbg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucHJvcHJpb19vYnMuc3Rk":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucmV3YXJkcy5tZWFu":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMucmV3YXJkcy5zdGQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMudmFsdWVzLm1lYW4=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c2NhbGVyLnNjYWxlcl9wYXJhbXMudmFsdWVzLnN0ZA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4wLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4xLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4yLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmF0dGVudGlvbl8yLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8zLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5LmxheWVycy4zLmxheWVybm9ybV8zLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X2RlYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X2RlYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cG9saWN5X3F1ZXJ5":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmF0dGVudGlvbi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmRlbnNlX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8xLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkLmxheWVycy4wLmxheWVybm9ybV8yLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2RlYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2RlYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","cmV3YXJkX2VtYi5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWUubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","dmFsdWVfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
\ No newline at end of file
diff --git a/dimladder/d1024/5000/default/array_metadatas/process_0 b/dimladder/d1024/5000/default/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..9a849e5ab26571fd7609f10b06d6312c4a32ed8a
--- /dev/null
+++ b/dimladder/d1024/5000/default/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.rngs.default.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.rngs.default.key", "write_shape": [2], "chunk_shape": [2], "ext_metadata": {"random_key_impl": "threefry2x32"}}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.bn2.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn1.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.mean", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.bn2.var", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn1.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.bn2.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_bn.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn1.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.mean", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.bn2.var", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn1.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.bn2.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_bn.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn1.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.mean", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.bn2.var", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn1.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.bn2.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_bn.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn1.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.mean", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.bn2.var", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.actions.mean", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.actions.std", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.proprio_obs.mean", "write_shape": [20], "chunk_shape": [20], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.proprio_obs.std", "write_shape": [20], "chunk_shape": [20], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.rewards.mean", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.rewards.std", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.values.mean", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "scaler.scaler_params.values.std", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/dimladder/d1024/5000/default/d/009e52348cb8d6fee0c34d716ee60bdf b/dimladder/d1024/5000/default/d/009e52348cb8d6fee0c34d716ee60bdf
new file mode 100644
index 0000000000000000000000000000000000000000..d08d42a851c41f941000b12fa0a94a279fffe419
Binary files /dev/null and b/dimladder/d1024/5000/default/d/009e52348cb8d6fee0c34d716ee60bdf differ
diff --git a/dimladder/d1024/5000/default/manifest.ocdbt b/dimladder/d1024/5000/default/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..94530b994cb56f60151a2323ddac424b858ef36a
Binary files /dev/null and b/dimladder/d1024/5000/default/manifest.ocdbt differ
diff --git a/dimladder/d1024/5000/default/ocdbt.process_0/d/151520133c9efda8afafbeedd4ae3547 b/dimladder/d1024/5000/default/ocdbt.process_0/d/151520133c9efda8afafbeedd4ae3547
new file mode 100644
index 0000000000000000000000000000000000000000..9657c84670b083d1c6d448948a884680895a647f
Binary files /dev/null and b/dimladder/d1024/5000/default/ocdbt.process_0/d/151520133c9efda8afafbeedd4ae3547 differ
diff --git a/dimladder/d1024/5000/default/ocdbt.process_0/d/3f4aca292c69f439aca31bc5f44b6c57 b/dimladder/d1024/5000/default/ocdbt.process_0/d/3f4aca292c69f439aca31bc5f44b6c57
new file mode 100644
index 0000000000000000000000000000000000000000..39a59cde6991688b27faba47140719c2d48d2007
Binary files /dev/null and b/dimladder/d1024/5000/default/ocdbt.process_0/d/3f4aca292c69f439aca31bc5f44b6c57 differ
diff --git a/dimladder/d1024/5000/default/ocdbt.process_0/d/43bb43789f770ca9e2f6da10d599eaa7 b/dimladder/d1024/5000/default/ocdbt.process_0/d/43bb43789f770ca9e2f6da10d599eaa7
new file mode 100644
index 0000000000000000000000000000000000000000..c1395f7b053338f12a8a208fb1f8f1058267ed79
--- /dev/null
+++ b/dimladder/d1024/5000/default/ocdbt.process_0/d/43bb43789f770ca9e2f6da10d599eaa7
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1f8a52dded4e7b0a9a83b28816d34ed8e6eea0fcccf6116d1021fed1bd93641
+size 399650816
diff --git a/dimladder/d1024/5000/default/ocdbt.process_0/d/6a2316abc2eccdc045bb5109d7564b50 b/dimladder/d1024/5000/default/ocdbt.process_0/d/6a2316abc2eccdc045bb5109d7564b50
new file mode 100644
index 0000000000000000000000000000000000000000..d3de95e53b1e0752a230ea395bc378ebecf3bfb2
Binary files /dev/null and b/dimladder/d1024/5000/default/ocdbt.process_0/d/6a2316abc2eccdc045bb5109d7564b50 differ
diff --git a/dimladder/d1024/5000/default/ocdbt.process_0/d/85c3547ded9ec94ce6e148c4d8551ff5 b/dimladder/d1024/5000/default/ocdbt.process_0/d/85c3547ded9ec94ce6e148c4d8551ff5
new file mode 100644
index 0000000000000000000000000000000000000000..7c051955298d7bbf225d51fb1b84855884ef3ad2
--- /dev/null
+++ b/dimladder/d1024/5000/default/ocdbt.process_0/d/85c3547ded9ec94ce6e148c4d8551ff5
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84306ca958132320456312a44b7a231c50079d801758656dea41641725d7534f
+size 135966720
diff --git a/dimladder/d1024/5000/default/ocdbt.process_0/d/9cedb10522084d797d68e95e4dadbbbf b/dimladder/d1024/5000/default/ocdbt.process_0/d/9cedb10522084d797d68e95e4dadbbbf
new file mode 100644
index 0000000000000000000000000000000000000000..465882bf6992569f157b41c73b5203a4ea126489
--- /dev/null
+++ b/dimladder/d1024/5000/default/ocdbt.process_0/d/9cedb10522084d797d68e95e4dadbbbf
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c6d4e695404f99b43dfcf5a670608b4cd4331476058fc4ddc221cb0c4621feb
+size 50507776
diff --git a/dimladder/d1024/5000/default/ocdbt.process_0/d/ff903252b14e2e6439bd63186259473a b/dimladder/d1024/5000/default/ocdbt.process_0/d/ff903252b14e2e6439bd63186259473a
new file mode 100644
index 0000000000000000000000000000000000000000..34f9c166653a4f5f5f67c2ff6410a0820c0bf4cb
--- /dev/null
+++ b/dimladder/d1024/5000/default/ocdbt.process_0/d/ff903252b14e2e6439bd63186259473a
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9949e60a445bd181ed4c36f1d344ff020580e6c12c8cc18edfd684837c079267
+size 50675712
diff --git a/dimladder/d1024/5000/default/ocdbt.process_0/manifest.ocdbt b/dimladder/d1024/5000/default/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..b5cd99cf9c5a109d73fafce1435946cf378707ab
Binary files /dev/null and b/dimladder/d1024/5000/default/ocdbt.process_0/manifest.ocdbt differ
diff --git a/dimladder/d1024/5000/opt_state/_METADATA b/dimladder/d1024/5000/opt_state/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..d2e91c3c1ba4762a96dba6939c0102b804fc9dac
--- /dev/null
+++ b/dimladder/d1024/5000/opt_state/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('opt_state', '1', '0', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'temporal_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('opt_state', '1', '0', 'mu', 'encoder', 'type_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'mu', 'fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('opt_state', '1', '0', 'mu', 'policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('opt_state', '1', '0', 'mu', 'policy_query')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'mu', 'reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('opt_state', '1', '0', 'mu', 'reward_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'mu', 'value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', '1', '0', 'mu', 'value_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4, 256]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'dynamics', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "dynamics", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'act_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'act_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "act_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 7, 3, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer1', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 64]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 64, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 64, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer2', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 128]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 128, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 128, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer3', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer3", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 256]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 256, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_bn', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_bn", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '0', 'downsample_conv', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "downsample_conv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1, 256, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'bn2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bn2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'image_enc', 'trunks', '0', 'layer4', '1', 'conv2', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "image_enc", "key_type": 2}, {"key": "trunks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layer4", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "conv2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 512, 512]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1556, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_mlp', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_mlp", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_norm', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'latent_norm', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 2048]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'proprio_obs_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "proprio_obs_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'temporal_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "temporal_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 1024]}}, "('opt_state', '1', '0', 'nu', 'encoder', 'type_enc')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "type_enc", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7, 1024]}}, "('opt_state', '1', '0', 'nu', 'fut_acts_embed', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "fut_acts_embed", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '0', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '1', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '2', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_1', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'attention_2', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "attention_2", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 8, 128]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_3', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy', 'layers', '3', 'layernorm_3', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "layernorm_3", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'policy_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('opt_state', '1', '0', 'nu', 'policy_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 7]}}, "('opt_state', '1', '0', 'nu', 'policy_query')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "policy_query", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'nu', 'reward_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1031, 1024]}}, "('opt_state', '1', '0', 'nu', 'reward_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "reward_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'key', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'key', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'out', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'out', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'query', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'query', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'value', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'attention', 'value', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attention", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 4096]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'dense_proj', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dense_proj", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4096, 1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_1', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_1', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_2', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value', 'layers', '0', 'layernorm_2', 'scale')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "layernorm_2", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value_dec', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('opt_state', '1', '0', 'nu', 'value_dec', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_dec", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'enc', 'output_layer', 'bias')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'enc', 'output_layer', 'kernel')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "enc", "key_type": 2}, {"key": "output_layer", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 1024]}}, "('opt_state', '1', '0', 'nu', 'value_emb', 'pos_emb')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "value_emb", "key_type": 2}, {"key": "pos_emb", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('opt_state', '1', '1', 'count')": {"key_metadata": [{"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('step',)": {"key_metadata": [{"key": "step", "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/dimladder/d1024/5000/opt_state/_sharding b/dimladder/d1024/5000/opt_state/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..17efc14e321f99e13537389be14be30a496b1ce3
--- /dev/null
+++ b/dimladder/d1024/5000/opt_state/_sharding
@@ -0,0 +1 @@
+{"b3B0X3N0YXRlLjEuMC5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS52YWx1ZV9lbWIucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9ub3JtLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLmxhdGVudF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnR5cGVfZW5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5lbmNvZGVyLnRlbXBvcmFsX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5mdXRfYWN0c19lbWJlZC5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5wb2xpY3lfcXVlcnk=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5tdS5yZXdhcmRfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5hdHRlbnRpb24udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZS5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9kZWMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIuZW5jLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS52YWx1ZV9lbWIucG9zX2VtYg==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMC5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMS5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5hdHRlbnRpb25fMi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5kZW5zZV9wcm9qLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMS5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMi5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5keW5hbWljcy5sYXllcnMuMi5sYXllcm5vcm1fMy5zY2FsZQ==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmFjdF9wcm9qLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjEuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjIuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjMuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2JuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMC5kb3duc2FtcGxlX2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5ibjIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmltYWdlX2VuYy50cnVua3MuMC5sYXllcjQuMS5jb252Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9ub3JtLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLmxhdGVudF9ub3JtLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnByb3ByaW9fb2JzX3Byb2oubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnR5cGVfZW5j":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5lbmNvZGVyLnRlbXBvcmFsX2VuYw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5lbmMub3V0cHV0X2xheWVyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5mdXRfYWN0c19lbWJlZC5wb3NfZW1i":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjAubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjEubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjIubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzEudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuYXR0ZW50aW9uXzIudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3kubGF5ZXJzLjMubGF5ZXJub3JtXzMuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5wb2xpY3lfcXVlcnk=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuYXR0ZW50aW9uLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAuZGVuc2VfcHJvai5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzEuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmQubGF5ZXJzLjAubGF5ZXJub3JtXzIuc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZGVjLm91dHB1dF9sYXllci5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLmVuYy5vdXRwdXRfbGF5ZXIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMC5udS5yZXdhcmRfZW1iLnBvc19lbWI=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","b3B0X3N0YXRlLjEuMS5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}","c3RlcA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"cuda:0\"}"}
\ No newline at end of file
diff --git a/dimladder/d1024/5000/opt_state/array_metadatas/process_0 b/dimladder/d1024/5000/opt_state/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..34e3283d0b868379783226fb713a9a90b9436271
--- /dev/null
+++ b/dimladder/d1024/5000/opt_state/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "opt_state.1.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.mu.value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_1.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.key.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.out.kernel", "write_shape": [4, 256, 1024], "chunk_shape": [4, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.query.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.bias", "write_shape": [4, 256], "chunk_shape": [4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.attention_2.value.kernel", "write_shape": [1024, 4, 256], "chunk_shape": [1024, 4, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.dynamics.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.act_proj.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.conv1.kernel", "write_shape": [7, 7, 3, 64], "chunk_shape": [7, 7, 3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.0.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.bn2.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv1.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer1.1.conv2.kernel", "write_shape": [3, 3, 64, 64], "chunk_shape": [3, 3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv1.kernel", "write_shape": [3, 3, 64, 128], "chunk_shape": [3, 3, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_bn.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.0.downsample_conv.kernel", "write_shape": [1, 1, 64, 128], "chunk_shape": [1, 1, 64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn1.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.bn2.scale", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv1.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer2.1.conv2.kernel", "write_shape": [3, 3, 128, 128], "chunk_shape": [3, 3, 128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv1.kernel", "write_shape": [3, 3, 128, 256], "chunk_shape": [3, 3, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_bn.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.0.downsample_conv.kernel", "write_shape": [1, 1, 128, 256], "chunk_shape": [1, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn1.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.bn2.scale", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv1.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer3.1.conv2.kernel", "write_shape": [3, 3, 256, 256], "chunk_shape": [3, 3, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv1.kernel", "write_shape": [3, 3, 256, 512], "chunk_shape": [3, 3, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_bn.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.0.downsample_conv.kernel", "write_shape": [1, 1, 256, 512], "chunk_shape": [1, 1, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.bn2.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv1.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.image_enc.trunks.0.layer4.1.conv2.kernel", "write_shape": [3, 3, 512, 512], "chunk_shape": [3, 3, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.0.kernel", "write_shape": [1556, 2048], "chunk_shape": [1556, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_mlp.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.latent_norm.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.0.kernel", "write_shape": [20, 2048], "chunk_shape": [20, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.layers.1.kernel", "write_shape": [2048, 2048], "chunk_shape": [2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.proprio_obs_proj.output_layer.kernel", "write_shape": [2048, 1024], "chunk_shape": [2048, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.temporal_enc", "write_shape": [6, 1024], "chunk_shape": [6, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.encoder.type_enc", "write_shape": [3, 1024], "chunk_shape": [3, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.enc.output_layer.kernel", "write_shape": [7, 1024], "chunk_shape": [7, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.fut_acts_embed.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.0.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.1.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.2.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_1.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.key.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.out.kernel", "write_shape": [8, 128, 1024], "chunk_shape": [8, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.query.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.bias", "write_shape": [8, 128], "chunk_shape": [8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.attention_2.value.kernel", "write_shape": [1024, 8, 128], "chunk_shape": [1024, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy.layers.3.layernorm_3.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.bias", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_dec.output_layer.kernel", "write_shape": [1024, 7], "chunk_shape": [1024, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.policy_query", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.enc.output_layer.kernel", "write_shape": [1031, 1024], "chunk_shape": [1031, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.reward_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.key.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.out.kernel", "write_shape": [1, 1024, 1024], "chunk_shape": [1, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.query.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.bias", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.attention.value.kernel", "write_shape": [1024, 1, 1024], "chunk_shape": [1024, 1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.layers.0.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.dense_proj.output_layer.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_1.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value.layers.0.layernorm_2.scale", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_dec.output_layer.kernel", "write_shape": [1024, 1], "chunk_shape": [1024, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.enc.output_layer.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.0.nu.value_emb.pos_emb", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "opt_state.1.1.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/dimladder/d1024/5000/opt_state/d/f0cb039b99129e426c79c6c9bad11ec9 b/dimladder/d1024/5000/opt_state/d/f0cb039b99129e426c79c6c9bad11ec9
new file mode 100644
index 0000000000000000000000000000000000000000..0f228861bf7f2beeefe33ab0d59a38a69f949e79
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/d/f0cb039b99129e426c79c6c9bad11ec9 differ
diff --git a/dimladder/d1024/5000/opt_state/manifest.ocdbt b/dimladder/d1024/5000/opt_state/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..879436b0ad55ec04df75650060f305f65c868d3f
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/manifest.ocdbt differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/0c7aaffb8a9bd2f5656e465a24030608 b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/0c7aaffb8a9bd2f5656e465a24030608
new file mode 100644
index 0000000000000000000000000000000000000000..bd9edad2868d3f10d934494464285c2f3fce7942
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/0c7aaffb8a9bd2f5656e465a24030608 differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/29a60f097c2a0152f81b8df34c752938 b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/29a60f097c2a0152f81b8df34c752938
new file mode 100644
index 0000000000000000000000000000000000000000..ad8eccb604d25b14d8c5c21d1893d56bd5882ee5
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/29a60f097c2a0152f81b8df34c752938 differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/3e211e8e65d994c5f09213377ba8a06c b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/3e211e8e65d994c5f09213377ba8a06c
new file mode 100644
index 0000000000000000000000000000000000000000..ad44c59da76f9c7ab56757140897f388bb531439
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/3e211e8e65d994c5f09213377ba8a06c differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/4526dfa97247fd77591b421289d1c1ba b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/4526dfa97247fd77591b421289d1c1ba
new file mode 100644
index 0000000000000000000000000000000000000000..5cad0396e923f7fd710a388db2f26cfd6c2d986e
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/4526dfa97247fd77591b421289d1c1ba differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/6e031a18b9dea7ed3c88f20a6609a891 b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/6e031a18b9dea7ed3c88f20a6609a891
new file mode 100644
index 0000000000000000000000000000000000000000..68819116d152d7ac5509821d2d00fe74d8195bce
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/6e031a18b9dea7ed3c88f20a6609a891 differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/7a4022fa8a810138d619aea82884f478 b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/7a4022fa8a810138d619aea82884f478
new file mode 100644
index 0000000000000000000000000000000000000000..0374c71251d03c59def6bf41c98c1ee485dceb0b
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/7a4022fa8a810138d619aea82884f478 differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/92c5678d992405b05ebcdc66fed25e61 b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/92c5678d992405b05ebcdc66fed25e61
new file mode 100644
index 0000000000000000000000000000000000000000..41f71228ff8ba61f5ab247abf2f1a5532a0f0ae9
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/92c5678d992405b05ebcdc66fed25e61 differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/a7d34fd8585833e4964ae3645c104f6a b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/a7d34fd8585833e4964ae3645c104f6a
new file mode 100644
index 0000000000000000000000000000000000000000..6e3ff2939f1cde9d620ae59dbca645d7104d14b4
--- /dev/null
+++ b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/a7d34fd8585833e4964ae3645c104f6a
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5884a24a3d6ba7d1a145422bd909ab170f4604c41282ff1da873ebab0939862a
+size 135916
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/b438955b4efbe0aaf8abf4b488aabfc0 b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/b438955b4efbe0aaf8abf4b488aabfc0
new file mode 100644
index 0000000000000000000000000000000000000000..a6150f50c130c456eb8bde4ec9040d4e8c6980fa
--- /dev/null
+++ b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/b438955b4efbe0aaf8abf4b488aabfc0
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a78fcca6754e2d537da9b3eb02ece2d32ae1c4105c56cc1cf4b4afc71f97bc24
+size 99565568
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/b620af38c89e747a4e29efa996879997 b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/b620af38c89e747a4e29efa996879997
new file mode 100644
index 0000000000000000000000000000000000000000..820619b5633a44f97ef11d38e5740e4f8c403068
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/b620af38c89e747a4e29efa996879997 differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/bc5c6aee1e3498353ed7403f9c7eea29 b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/bc5c6aee1e3498353ed7403f9c7eea29
new file mode 100644
index 0000000000000000000000000000000000000000..9b76b703138d2b7eeed39dded04456eef756722c
--- /dev/null
+++ b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/bc5c6aee1e3498353ed7403f9c7eea29
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3120b74f201e601d3b9d7874526bdce5b93143b596ff755c95164ef8ee5302d6
+size 1099460608
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/c3c5e66594e0899f80de113ec42c18a2 b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/c3c5e66594e0899f80de113ec42c18a2
new file mode 100644
index 0000000000000000000000000000000000000000..23561a7cbd154d5368c3b774c957c3ad547b6e69
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/c3c5e66594e0899f80de113ec42c18a2 differ
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/f7b7a396546af89363a458e1d243477e b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/f7b7a396546af89363a458e1d243477e
new file mode 100644
index 0000000000000000000000000000000000000000..b83b41284de944b5f444f146a37b98e402206bbd
--- /dev/null
+++ b/dimladder/d1024/5000/opt_state/ocdbt.process_0/d/f7b7a396546af89363a458e1d243477e
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dfa2add81a928fa4414438541d08465b14fd58967cfe8119c37045375809dea
+size 61059072
diff --git a/dimladder/d1024/5000/opt_state/ocdbt.process_0/manifest.ocdbt b/dimladder/d1024/5000/opt_state/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..c8ef7f09cdc9f3a3bfbc1ea090733dadb98624d4
Binary files /dev/null and b/dimladder/d1024/5000/opt_state/ocdbt.process_0/manifest.ocdbt differ
diff --git a/dimladder/d1024/diagnostics/step_10000/decoder_probe.npz b/dimladder/d1024/diagnostics/step_10000/decoder_probe.npz
new file mode 100644
index 0000000000000000000000000000000000000000..55c8ee22b80fae3ec95445aa320d89c2687d4dcd
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_10000/decoder_probe.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e643f8507e929944402472dc460ea7c955c869ec6cd1c4e3f87b909d9c291fb
+size 2810851
diff --git a/dimladder/d1024/diagnostics/step_10000/decoder_probe.png b/dimladder/d1024/diagnostics/step_10000/decoder_probe.png
new file mode 100644
index 0000000000000000000000000000000000000000..2ac26d026683e71ee955ae170c883604771fe812
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_10000/decoder_probe.png differ
diff --git a/dimladder/d1024/diagnostics/step_10000/latent_rollout_error.npz b/dimladder/d1024/diagnostics/step_10000/latent_rollout_error.npz
new file mode 100644
index 0000000000000000000000000000000000000000..93445b38b58479455084339a90552a94ade12799
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_10000/latent_rollout_error.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd160450989632ee37a4bc2eea3f31d3b841eb41272d8fb0203c4dfcde22f2d0
+size 467
diff --git a/dimladder/d1024/diagnostics/step_10000/latent_rollout_error.png b/dimladder/d1024/diagnostics/step_10000/latent_rollout_error.png
new file mode 100644
index 0000000000000000000000000000000000000000..7953209c6e23f0519733c66a2a6cf26a3b04cc56
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_10000/latent_rollout_error.png differ
diff --git a/dimladder/d1024/diagnostics/step_10000/latent_scatter.npz b/dimladder/d1024/diagnostics/step_10000/latent_scatter.npz
new file mode 100644
index 0000000000000000000000000000000000000000..9ef84c298153ef3950bc7be63c973f537ad2cfe1
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_10000/latent_scatter.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd8a1b8980d81907fe8ec62639287dc5c5324b25f5725d5d6764e317cfa0eb2b
+size 75079
diff --git a/dimladder/d1024/diagnostics/step_10000/latent_scatter.png b/dimladder/d1024/diagnostics/step_10000/latent_scatter.png
new file mode 100644
index 0000000000000000000000000000000000000000..776d853eeb9026e311f09f9d0a75cf939f89e737
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_10000/latent_scatter.png differ
diff --git a/dimladder/d1024/diagnostics/step_10000/reward_value_curves.npz b/dimladder/d1024/diagnostics/step_10000/reward_value_curves.npz
new file mode 100644
index 0000000000000000000000000000000000000000..07741db7d969d69752799da14062c74da63de2df
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_10000/reward_value_curves.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccca75f705fc317cf1a2b83f911b4a321647ba39dd1d60f6b93abb806a4b4d31
+size 59221
diff --git a/dimladder/d1024/diagnostics/step_10000/reward_value_curves.png b/dimladder/d1024/diagnostics/step_10000/reward_value_curves.png
new file mode 100644
index 0000000000000000000000000000000000000000..ae3c9d462f70a4a802dcfb160aabf0fe27a9282b
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_10000/reward_value_curves.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:169ca30c2b6f52e5cb1904723600fee6534e68f189279b473deb5802b02b61b3
+size 771398
diff --git a/dimladder/d1024/diagnostics/step_12500/decoder_probe.npz b/dimladder/d1024/diagnostics/step_12500/decoder_probe.npz
new file mode 100644
index 0000000000000000000000000000000000000000..814e37981d54c1a577568e2c5011f7cfeb386781
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_12500/decoder_probe.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b82dd87d6eb8807c82eb9ffb7da1868b3e8e807ae6bf63d6cbaff6ca49ab489
+size 2808659
diff --git a/dimladder/d1024/diagnostics/step_12500/decoder_probe.png b/dimladder/d1024/diagnostics/step_12500/decoder_probe.png
new file mode 100644
index 0000000000000000000000000000000000000000..887de33ea13ccb23409a6a58b9d802adc23c62ed
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_12500/decoder_probe.png differ
diff --git a/dimladder/d1024/diagnostics/step_12500/latent_rollout_error.npz b/dimladder/d1024/diagnostics/step_12500/latent_rollout_error.npz
new file mode 100644
index 0000000000000000000000000000000000000000..5cd842ffdae2ecb8e95b755589ddd3fa98180026
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_12500/latent_rollout_error.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:40d59a25c338aab12478dd329f3bed7d5d5f39be1a87975e61f177397516e7a2
+size 467
diff --git a/dimladder/d1024/diagnostics/step_12500/latent_rollout_error.png b/dimladder/d1024/diagnostics/step_12500/latent_rollout_error.png
new file mode 100644
index 0000000000000000000000000000000000000000..5964a409a1b41ea11511c68cb09d85f413672102
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_12500/latent_rollout_error.png differ
diff --git a/dimladder/d1024/diagnostics/step_12500/latent_scatter.npz b/dimladder/d1024/diagnostics/step_12500/latent_scatter.npz
new file mode 100644
index 0000000000000000000000000000000000000000..ce52cbce7e50d43f2bb42f7bffd92c362c33a2f9
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_12500/latent_scatter.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09396ed1442da9cdb613ebac236d71b6985f7776309b50f99012a6048809239a
+size 75164
diff --git a/dimladder/d1024/diagnostics/step_12500/latent_scatter.png b/dimladder/d1024/diagnostics/step_12500/latent_scatter.png
new file mode 100644
index 0000000000000000000000000000000000000000..893fda5b93d6d336cf550526b87b88c285640332
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_12500/latent_scatter.png differ
diff --git a/dimladder/d1024/diagnostics/step_12500/reward_value_curves.npz b/dimladder/d1024/diagnostics/step_12500/reward_value_curves.npz
new file mode 100644
index 0000000000000000000000000000000000000000..9c44de12bb44b6eaa9c4bfb350247de26a804a76
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_12500/reward_value_curves.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9854f5245c93e10c877b7633f123d72473e3dc340050f48385e3218dd42907a
+size 59274
diff --git a/dimladder/d1024/diagnostics/step_12500/reward_value_curves.png b/dimladder/d1024/diagnostics/step_12500/reward_value_curves.png
new file mode 100644
index 0000000000000000000000000000000000000000..f0888e4ec4a6ac77f5b8934e68d212b1e28ea003
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_12500/reward_value_curves.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fea4f09eb18df5d7ef16cc2d967f7542a78fedf176d36dfc7d5020e3b6e70a35
+size 757578
diff --git a/dimladder/d1024/diagnostics/step_15000/decoder_probe.npz b/dimladder/d1024/diagnostics/step_15000/decoder_probe.npz
new file mode 100644
index 0000000000000000000000000000000000000000..6ae02809302e24e7325bdfa5c9af185f563a0be1
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_15000/decoder_probe.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b8c2313f0b08994f0fb052a9e235882a72dfebd756cb32b992c7cabcd9ea0a5
+size 2893452
diff --git a/dimladder/d1024/diagnostics/step_15000/decoder_probe.png b/dimladder/d1024/diagnostics/step_15000/decoder_probe.png
new file mode 100644
index 0000000000000000000000000000000000000000..b0e8582909608ae3d40de3c71a64181657c7094c
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_15000/decoder_probe.png differ
diff --git a/dimladder/d1024/diagnostics/step_15000/latent_rollout_error.npz b/dimladder/d1024/diagnostics/step_15000/latent_rollout_error.npz
new file mode 100644
index 0000000000000000000000000000000000000000..eca3628cd32629aad881681a9b96c0b6ba7c4502
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_15000/latent_rollout_error.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3418e535212246503584991eef76adac40a474a51b56e621945b6d3c2b181d72
+size 467
diff --git a/dimladder/d1024/diagnostics/step_15000/latent_rollout_error.png b/dimladder/d1024/diagnostics/step_15000/latent_rollout_error.png
new file mode 100644
index 0000000000000000000000000000000000000000..14a0c6c0021b767ab999b071b21ac8ca4c9b5ae9
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_15000/latent_rollout_error.png differ
diff --git a/dimladder/d1024/diagnostics/step_15000/latent_scatter.npz b/dimladder/d1024/diagnostics/step_15000/latent_scatter.npz
new file mode 100644
index 0000000000000000000000000000000000000000..75ca56bbca3e40565f0e8da5644356717a4db614
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_15000/latent_scatter.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c31e4a1bbca55352a5df1b2f2d62239769d55a2aa55eb96186abf5a97a74a0a
+size 75203
diff --git a/dimladder/d1024/diagnostics/step_15000/latent_scatter.png b/dimladder/d1024/diagnostics/step_15000/latent_scatter.png
new file mode 100644
index 0000000000000000000000000000000000000000..ec6b81d2ad037cc420ceccd3843a0ba56754c8e0
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_15000/latent_scatter.png differ
diff --git a/dimladder/d1024/diagnostics/step_15000/reward_value_curves.npz b/dimladder/d1024/diagnostics/step_15000/reward_value_curves.npz
new file mode 100644
index 0000000000000000000000000000000000000000..86061c7a974eb856a7772dd9bf6b25951af00026
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_15000/reward_value_curves.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e506883d31bc276ffd9c39406fa16e871bb4995ce87a2eeea601f7c0b94b3c0b
+size 59304
diff --git a/dimladder/d1024/diagnostics/step_15000/reward_value_curves.png b/dimladder/d1024/diagnostics/step_15000/reward_value_curves.png
new file mode 100644
index 0000000000000000000000000000000000000000..a1b105063918f7ab557eed690cb293a5ed7978a5
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_15000/reward_value_curves.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6fcf0e54808f5454c0371ceb0983cb89231cf9a002d0e1d3da4ecd141d0820b
+size 782506
diff --git a/dimladder/d1024/diagnostics/step_5000/decoder_probe.npz b/dimladder/d1024/diagnostics/step_5000/decoder_probe.npz
new file mode 100644
index 0000000000000000000000000000000000000000..be18ad2eb5dd28a1ab163e46afb98ab23e6f9ce8
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_5000/decoder_probe.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:35266ee1c83b310572d72911f9502b7e6a6d8fa27bc7300e2e2e64776d28b60c
+size 2805752
diff --git a/dimladder/d1024/diagnostics/step_5000/decoder_probe.png b/dimladder/d1024/diagnostics/step_5000/decoder_probe.png
new file mode 100644
index 0000000000000000000000000000000000000000..9f8196d511f7bc0a1fd41517e965b01408ad710c
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_5000/decoder_probe.png differ
diff --git a/dimladder/d1024/diagnostics/step_5000/latent_rollout_error.npz b/dimladder/d1024/diagnostics/step_5000/latent_rollout_error.npz
new file mode 100644
index 0000000000000000000000000000000000000000..5db347b1b908cf1a9f85aec3779da55bae947008
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_5000/latent_rollout_error.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b8ff1d66672db9efc7e4d8da5796c6b21fb6b789511b085273a72b78e6f5edb
+size 467
diff --git a/dimladder/d1024/diagnostics/step_5000/latent_rollout_error.png b/dimladder/d1024/diagnostics/step_5000/latent_rollout_error.png
new file mode 100644
index 0000000000000000000000000000000000000000..0232f2dbaf30d82e70d744a43c4a19f90d10f53b
Binary files /dev/null and b/dimladder/d1024/diagnostics/step_5000/latent_rollout_error.png differ
diff --git a/dimladder/d1024/diagnostics/step_5000/latent_scatter.npz b/dimladder/d1024/diagnostics/step_5000/latent_scatter.npz
new file mode 100644
index 0000000000000000000000000000000000000000..e3427b7712c4a1d452cbe7fcad29a074149e8d34
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_5000/latent_scatter.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf1e03a6cffd54cc904ed1b6c45673c38ff5142d1b49c8e84dc790b00204d1a8
+size 75101
diff --git a/dimladder/d1024/diagnostics/step_5000/latent_scatter.png b/dimladder/d1024/diagnostics/step_5000/latent_scatter.png
new file mode 100644
index 0000000000000000000000000000000000000000..b905c44ca76a83375dbcf8b4e8578aec6c738e70
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_5000/latent_scatter.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16a3e3a41e44ca09c88b9d4c470c24c701988fe17191b242d3054698867e2544
+size 111371
diff --git a/dimladder/d1024/diagnostics/step_5000/reward_value_curves.npz b/dimladder/d1024/diagnostics/step_5000/reward_value_curves.npz
new file mode 100644
index 0000000000000000000000000000000000000000..29f3856493afb116a33cead3096c50053188064d
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_5000/reward_value_curves.npz
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f61a397e279b0a0b5f448d41639a26a3fe07c34be5564c1b4a9001510ae12507
+size 59695
diff --git a/dimladder/d1024/diagnostics/step_5000/reward_value_curves.png b/dimladder/d1024/diagnostics/step_5000/reward_value_curves.png
new file mode 100644
index 0000000000000000000000000000000000000000..41fa8a54a6765470ec79d1acef054025e0aaf29a
--- /dev/null
+++ b/dimladder/d1024/diagnostics/step_5000/reward_value_curves.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1afbc0e582874e2164ec97aa47895435fae528095a3212391d52d10e62628c0e
+size 818563
diff --git a/dimladder/d1024/model_config.yaml b/dimladder/d1024/model_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..bd85101d02559f415dfb957384d2e83ad98aa636
--- /dev/null
+++ b/dimladder/d1024/model_config.yaml
@@ -0,0 +1,220 @@
+checkpoint:
+  enabled: true
+  max_to_keep: null
+  resume_checkpoint: null
+data:
+  dataset_name: simdist_nodr
+  num_test_workers: 4
+  num_train_workers: 8
+  paper_scale_steps: 25000000
+  prefetch_factor: 2
+eval:
+  decoder_probe:
+    enabled: true
+    hidden_size: 128
+    image:
+      base_channels: 128
+      batch: 32
+      enabled: true
+      learning_rate: 0.001
+    learning_rate: 0.001
+    steps_per_eval: 200
+  diagnostics:
+    dirname: diagnostics
+    enabled: true
+  episodes:
+    batch_size: 32
+    curve_stride: 5
+    enabled: true
+    holdout: 128
+    match_frac: 0.5
+    num_auc_episodes: 512
+    num_curve_episodes: 16
+    success:
+      metric: episode_return
+      threshold: null
+  max_batches: 8
+  scatter_points: 5000
+heads: {}
+loss:
+  type: world_model
+  weights:
+    action: 4.0
+    latent_dynamics: 1.0
+    reward: 1.0
+    value: 1.0
+model:
+  dataset:
+    augmentations:
+      add_noise: true
+      image:
+        color_jitter:
+          brightness: 0.2
+          contrast: 0.2
+          enabled: true
+          hue: 0.02
+          saturation: 0.2
+        enabled: true
+        gaussian_blur:
+          enabled: true
+          kernel_size: 5
+          prob: 0.5
+          sigma:
+          - 0.1
+          - 1.5
+        random_crop:
+          enabled: true
+          padding:
+          - 4
+          - 6
+    history_length: 5
+    jpeg_backend: auto
+    prediction_length: 5
+    type: manipulator_world_model
+  dropout:
+    attention: 0.05
+    mlp: 0.2
+  dynamics:
+    attention:
+      heads: 4
+      layers: 3
+      mask: causal
+      mlp_hidden_size_factor: 4
+  embedding:
+    future_acts_layers: 0
+    future_cmds_layers: 0
+    mlp_hidden_size_factor: 2
+    reward_layers: 0
+    value_layers: 0
+  encoder:
+    action_layers: 0
+    extero_obs:
+      image_resnet:
+        freeze_batch_norm: true
+        normalize_input: true
+        pretrained: true
+        share_trunk: true
+    latent_layers: 2
+    latent_norm: true
+    mlp_hidden_size_factor: 2
+    proprio_obs_layers: 2
+  latent_dim: 1024
+  policy:
+    attention:
+      heads: 8
+      layers: 4
+      mask: causal
+      mlp_hidden_size_factor: 4
+    decoder:
+      layers: 0
+      mlp_hidden_size_factor: 2
+  reward:
+    attention:
+      heads: 1
+      layers: 1
+      mask: null
+      mlp_hidden_size_factor: 4
+    decoder:
+      layers: 0
+      mlp_hidden_size_factor: 2
+  type: manipulator_world_model
+  value:
+    attention:
+      heads: 1
+      layers: 1
+      mask: null
+      mlp_hidden_size_factor: 4
+    decoder:
+      layers: 0
+      mlp_hidden_size_factor: 2
+run_name: nodr_d1024_r2
+scaler_params_struct:
+  actions: 7
+  commands: 0
+  extero_obs: 0
+  proprio_obs: 20
+  rewards: 1
+  values: 1
+system:
+  actions:
+  - max_noise: 1.0
+    min_noise: 0.05
+    name: ee_delta_x
+  - max_noise: 1.0
+    min_noise: 0.05
+    name: ee_delta_y
+  - max_noise: 1.0
+    min_noise: 0.05
+    name: ee_delta_z
+  - max_noise: 1.0
+    min_noise: 0.05
+    name: ee_delta_rx
+  - max_noise: 1.0
+    min_noise: 0.05
+    name: ee_delta_ry
+  - max_noise: 1.0
+    min_noise: 0.05
+    name: ee_delta_rz
+  - flip_prob_max: 0.1
+    flip_prob_min: 0.01
+    max_noise: 0.0
+    min_noise: 0.0
+    name: gripper_binary
+  cmd:
+    dim: 0
+  extero_obs:
+    types:
+    - dim:
+      - 120
+      - 160
+      - 3
+      name: front_rgb
+      noise: 0.01
+    - dim:
+      - 120
+      - 160
+      - 3
+      name: side_rgb
+      noise: 0.01
+    - dim:
+      - 120
+      - 160
+      - 3
+      name: wrist_rgb
+      noise: 0.01
+  name: omnireset_ur10e
+  proprio_obs:
+    types:
+    - dim: 1
+      name: last_gripper_action
+      noise: 0.01
+    - dim: 6
+      name: last_arm_action
+      noise: 0.01
+    - dim: 6
+      name: arm_joint_pos
+      noise: 0.005
+    - dim: 6
+      name: end_effector_pose
+      noise: 0.003
+    - dim: 1
+      name: binary_contact
+      noise: 0.05
+training:
+  batch_size: 512
+  decay_steps: 400000
+  end_learning_rate: 0.0001
+  eval_interval: 5000
+  grad_clip_norm: 1.0
+  learning_rate: 0.0002
+  max_steps: 15000
+  num_epochs: 100
+  seed: 42
+  training_data_ratio: 0.98
+  warmup_steps: 10000
+wandb:
+  dir: null
+  entity: null
+  log: false
+  mode: online
+  project: simdist_train_model