diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..1f79f27b854d111bb55aa418bcaa82861180127e 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,30 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+permanent/step_200000/checkpoint_200000/ocdbt.process_0/d/b2740f6c6449f1c164c9bc7918fe176e filter=lfs diff=lfs merge=lfs -text
+opt_state/checkpoint_400000/ocdbt.process_0/d/4e81829b310436056edf3761802dc154 filter=lfs diff=lfs merge=lfs -text
+permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/3e028d2bbb4b009cc0a1652900d6dbae filter=lfs diff=lfs merge=lfs -text
+latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/9a268f7c8b2e0c77a42b67cd4920356f filter=lfs diff=lfs merge=lfs -text
+checkpoint_400000/ocdbt.process_0/d/b137f1e9ebe14cb480bc3fa3fcfd6e1a filter=lfs diff=lfs merge=lfs -text
+latest/opt_state/checkpoint_400000/ocdbt.process_0/d/55f24edf08ee69029b9ece4618cc887a filter=lfs diff=lfs merge=lfs -text
+latest/ema/checkpoint_400000/ocdbt.process_0/d/80d4d7186e25094c2a630a6d91f24a8f filter=lfs diff=lfs merge=lfs -text
+permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/4535724fe7708c015e5ce5eb473d37a7 filter=lfs diff=lfs merge=lfs -text
+latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/1e1dbbbd4fa63924d325bb36de9f1789 filter=lfs diff=lfs merge=lfs -text
+permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/380c55ccd750470f7dadc2175c156a5b filter=lfs diff=lfs merge=lfs -text
+permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/99be9958200bf1b41a1dbbc186832b4d filter=lfs diff=lfs merge=lfs -text
+checkpoint_400000/ocdbt.process_0/d/197494707a6ad143f425e535bbb35d03 filter=lfs diff=lfs merge=lfs -text
+permanent/step_100000/ema/checkpoint_100000/ocdbt.process_0/d/5cdb2bacecd19679b6575ae33490dd36 filter=lfs diff=lfs merge=lfs -text
+permanent/step_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/b399e969b4326f5b52f8e76b7cc2ae7c filter=lfs diff=lfs merge=lfs -text
+latest/ema/checkpoint_400000/ocdbt.process_0/d/1de9dc33f844966b10fa8a34bb366867 filter=lfs diff=lfs merge=lfs -text
+permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/d7fd7d3fdd1c351adac17ed11a0323e3 filter=lfs diff=lfs merge=lfs -text
+permanent/step_400000/opt_state/checkpoint_400000/ocdbt.process_0/d/b6c8e12debdd19d0beeece616d43dcf5 filter=lfs diff=lfs merge=lfs -text
+permanent/step_400000/opt_state/checkpoint_400000/ocdbt.process_0/d/d53aa29264c20ee8d87912444fba44d1 filter=lfs diff=lfs merge=lfs -text
+permanent/step_200000/opt_state/checkpoint_200000/ocdbt.process_0/d/c827e39babc608b717bc851428d44939 filter=lfs diff=lfs merge=lfs -text
+permanent/step_200000/ema/checkpoint_200000/ocdbt.process_0/d/95af092583d380070ceacbee12c385fd filter=lfs diff=lfs merge=lfs -text
+permanent/step_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/4c4188133a37300d5ada5ac8a3105934 filter=lfs diff=lfs merge=lfs -text
+latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/ca22174a382c5e795a946cf230562125 filter=lfs diff=lfs merge=lfs -text
+permanent/step_100000/opt_state/checkpoint_100000/ocdbt.process_0/d/67d4789f8c2b945a3abb6a69f2311204 filter=lfs diff=lfs merge=lfs -text
+permanent/step_100000/checkpoint_100000/ocdbt.process_0/d/ecf6faa66144dab60e1504644745b153 filter=lfs diff=lfs merge=lfs -text
+permanent/step_400000/checkpoint_400000/ocdbt.process_0/d/058cbf30ed1dce8e04ea12f92cc3522e filter=lfs diff=lfs merge=lfs -text
+permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/24733009192379f8f59f98e1c9ddf6ad filter=lfs diff=lfs merge=lfs -text
+permanent/step_400000/ema/checkpoint_400000/ocdbt.process_0/d/910e1420ad62226535ee916a97e440f5 filter=lfs diff=lfs merge=lfs -text
diff --git a/checkpoint_400000/_CHECKPOINT_METADATA b/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..862e57576a39ab149aebfe62019f92c65ec6a55d
--- /dev/null
+++ b/checkpoint_400000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778652475406833336, "commit_timestamp_nsecs": 1778652491174548258, "custom_metadata": {}}
\ No newline at end of file
diff --git a/checkpoint_400000/_METADATA b/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..bbd0485d06d3a08df040be11c549f90fb114ae53
--- /dev/null
+++ b/checkpoint_400000/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('backbone', 'DiTBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_10', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_10', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_11', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_11', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_12', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_12', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_12', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_12', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_13', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_13', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_13', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_13', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_14', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_14', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_14', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_14', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_15', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_15', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_15', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_15', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_16', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_16', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_16', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_16', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_17', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_17', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_17', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_17', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_18', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_18', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_18', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_18', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_19', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_19', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_19', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_19', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_2', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_2', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_20', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_20', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_20', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_20', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_21', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_21', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_21', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_21', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_22', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_22', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_22', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_22', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_22', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_22', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_23', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_23', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_23', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_23', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_23', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_23', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_24', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_24', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_24', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_24', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_24', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_24', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_25', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_25', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_25', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_25', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_25', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_25', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_26', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_26', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_26', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_26', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_26', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_26', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_27', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_27', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_27', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_27', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_27', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_27', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_3', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_3', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_4', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_4', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_5', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_5', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_6', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_6', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_7', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_7', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_8', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_8', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_9', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_9', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'FinalLayer_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('backbone', 'FinalLayer_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('backbone', 'FinalLayer_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304]}}, "('backbone', 'FinalLayer_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 2304]}}, "('backbone', 'LabelEmbedder_0', 'Embed_0', 'embedding')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "LabelEmbedder_0", "key_type": 2}, {"key": "Embed_0", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('backbone', 'PatchedPatchEmbed_0', 'proj', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'PatchedPatchEmbed_0', 'proj', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 1152]}}, "('backbone', 'TimestepEmbedder_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'TimestepEmbedder_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1152]}}, "('backbone', 'TimestepEmbedder_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'TimestepEmbedder_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 1152]}}, "('predictor', 'blocks_0', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_0', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_0', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_0', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_0', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_0', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_0', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_0', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_0', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_0', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_0', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_0', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_1', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_1', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_1', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_1', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_1', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_1', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_1', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_1', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_1', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_1', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_1', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_1', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_10', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_10', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_10', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_10', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_10', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_10', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_10', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_10', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_10', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_10', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_10', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_10', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_11', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_11', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_11', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_11', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('predictor', 'blocks_11', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_11', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_11', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_11', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_11', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_11', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_11', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_11', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_11', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_11', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_11', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_11', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_11', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_11', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_11', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_11', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_11', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_11', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_2', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_2', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_2', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_2', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_2', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_2', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_2', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_2', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_2', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_2', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_2', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_2', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_3', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_3', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_3', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_3', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('predictor', 'blocks_3', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_3', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_3', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_3', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_3', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_3', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_3', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_3', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_3', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_3', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_3', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_3', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_3', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_3', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_3', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_3', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_3', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_3', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_4', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_4', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_4', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_4', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_4', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_4', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_4', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_4', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_4', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_4', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_4', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_4', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_5', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_5', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_5', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_5', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_5', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_5', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_5', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_5', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_5', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_5', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_5', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_5', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_6', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_6', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_6', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_6', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_6', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_6', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_6', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_6', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_6', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_6', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_6', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_6', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_7', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_7', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_7', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_7', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('predictor', 'blocks_7', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_7', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_7', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_7', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_7', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_7', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_7', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_7', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_7', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_7', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_7', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_7', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_7', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_7', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_7', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_7', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_7', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_7', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_8', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_8', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_8', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_8', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_8', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_8', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_8', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_8', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_8', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_8', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_8', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_8', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_9', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_9', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_9', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_9', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_9', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_9', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_9', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_9', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_9', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_9', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_9', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_9', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'class_input_embed', 'embedding')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "class_input_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('predictor', 'cond_delta_embed', 'embedding')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_delta_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [28, 32]}}, "('predictor', 'cond_layer_embed', 'embedding')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_layer_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [29, 32]}}, "('predictor', 'cond_out', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('predictor', 'cond_out', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('predictor', 'cond_t_proj', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('predictor', 'cond_t_proj', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('predictor', 'final_ln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'final_ln', 'scale')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'gamma_out')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "gamma_out", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('predictor', 'in_proj', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'in_proj', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 480]}}, "('predictor', 'mag_head', 'adaln_mod', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [964]}}, "('predictor', 'mag_head', 'adaln_mod', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 964]}}, "('predictor', 'mag_head', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482]}}, "('predictor', 'mag_head', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 482]}}, "('predictor', 'mag_head', 'pw1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120]}}, "('predictor', 'mag_head', 'pw1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482, 120]}}, "('predictor', 'mag_head', 'pw2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('predictor', 'mag_head', 'pw2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120, 1]}}, "('predictor', 'out_proj', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('predictor', 'out_proj', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1152]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/checkpoint_400000/_sharding b/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..384403235a88279737437bfee68391867a2b8108
--- /dev/null
+++ b/checkpoint_400000/_sharding
@@ -0,0 +1 @@
+{"YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRmluYWxMYXllcl8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRmluYWxMYXllcl8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRmluYWxMYXllcl8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRmluYWxMYXllcl8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuTGFiZWxFbWJlZGRlcl8wLkVtYmVkXzAuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuUGF0Y2hlZFBhdGNoRW1iZWRfMC5wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuUGF0Y2hlZFBhdGNoRW1iZWRfMC5wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuVGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuVGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuVGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuVGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLmFkYWxuX21vZC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLmFkYWxuX21vZC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLnB3MS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLnB3MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLnB3Mi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLnB3Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm91dF9wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm91dF9wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5jb252X2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5jb252X2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5kd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5kd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5wd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5wd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5jb252X2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5jb252X2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5kd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5kd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5wd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5wd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNsYXNzX2lucHV0X2VtYmVkLmVtYmVkZGluZw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfZGVsdGFfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfb3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfb3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfbGF5ZXJfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfdF9wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfdF9wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmZpbmFsX2xuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmZpbmFsX2xuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmdhbW1hX291dA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmluX3Byb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmluX3Byb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}"}
\ No newline at end of file
diff --git a/checkpoint_400000/array_metadatas/process_0 b/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..d267d96a4c3b8ebf13b6f6b35fac282dadea1079
--- /dev/null
+++ b/checkpoint_400000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_0.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_1.bias", "write_shape": [2304], "chunk_shape": [2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_1.kernel", "write_shape": [1152, 2304], "chunk_shape": [1152, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.LabelEmbedder_0.Embed_0.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.PatchedPatchEmbed_0.proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.PatchedPatchEmbed_0.proj.kernel", "write_shape": [16, 1152], "chunk_shape": [16, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_0.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_0.kernel", "write_shape": [256, 1152], "chunk_shape": [256, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_1.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_1.kernel", "write_shape": [1152, 1152], "chunk_shape": [1152, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.class_input_embed.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_delta_embed.embedding", "write_shape": [28, 32], "chunk_shape": [28, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_layer_embed.embedding", "write_shape": [29, 32], "chunk_shape": [29, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_out.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_out.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_t_proj.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_t_proj.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.final_ln.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.final_ln.scale", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.gamma_out", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.in_proj.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.in_proj.kernel", "write_shape": [1152, 480], "chunk_shape": [1152, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.adaln_mod.bias", "write_shape": [964], "chunk_shape": [964], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.adaln_mod.kernel", "write_shape": [32, 964], "chunk_shape": [32, 964], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.dwconv.bias", "write_shape": [482], "chunk_shape": [482], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.dwconv.kernel", "write_shape": [3, 3, 1, 482], "chunk_shape": [3, 3, 1, 482], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.pw1.bias", "write_shape": [120], "chunk_shape": [120], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.pw1.kernel", "write_shape": [482, 120], "chunk_shape": [482, 120], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.pw2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.pw2.kernel", "write_shape": [120, 1], "chunk_shape": [120, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.out_proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.out_proj.kernel", "write_shape": [480, 1152], "chunk_shape": [480, 1152], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/checkpoint_400000/d/89ccddba872ec98f9f7196e91e61adce b/checkpoint_400000/d/89ccddba872ec98f9f7196e91e61adce
new file mode 100644
index 0000000000000000000000000000000000000000..3f3ad70a5d57210936a19376a113d9239d80ca14
Binary files /dev/null and b/checkpoint_400000/d/89ccddba872ec98f9f7196e91e61adce differ
diff --git a/checkpoint_400000/manifest.ocdbt b/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..1b7cd64e9ceccb6f053536ceac377785abb6b640
Binary files /dev/null and b/checkpoint_400000/manifest.ocdbt differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/197494707a6ad143f425e535bbb35d03 b/checkpoint_400000/ocdbt.process_0/d/197494707a6ad143f425e535bbb35d03
new file mode 100644
index 0000000000000000000000000000000000000000..35fe017fa35120dbd51b52e6a47cbdd4303efaaa
--- /dev/null
+++ b/checkpoint_400000/ocdbt.process_0/d/197494707a6ad143f425e535bbb35d03
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb8dbd6bb973040f85b4820e24b26684864582012f77742ac735ca3552552fe0
+size 29597696
diff --git a/checkpoint_400000/ocdbt.process_0/d/391ddebdf84ae2ebc892e17e78d2e73f b/checkpoint_400000/ocdbt.process_0/d/391ddebdf84ae2ebc892e17e78d2e73f
new file mode 100644
index 0000000000000000000000000000000000000000..e9c0fc2831534687cb487bbffa7cb4e26b32033d
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/391ddebdf84ae2ebc892e17e78d2e73f differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/430e231c749c8133a2442b29fc42f886 b/checkpoint_400000/ocdbt.process_0/d/430e231c749c8133a2442b29fc42f886
new file mode 100644
index 0000000000000000000000000000000000000000..ba791b5bed403fdc40da332c7903ef4c06a7c5d7
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/430e231c749c8133a2442b29fc42f886 differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/4f6ad7f5f74533a1a9074df166e2e1ad b/checkpoint_400000/ocdbt.process_0/d/4f6ad7f5f74533a1a9074df166e2e1ad
new file mode 100644
index 0000000000000000000000000000000000000000..c51c0d64ba91b59851b521c102c4652027bae48c
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/4f6ad7f5f74533a1a9074df166e2e1ad differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/658f2dccd9887c974b4f574fe0511442 b/checkpoint_400000/ocdbt.process_0/d/658f2dccd9887c974b4f574fe0511442
new file mode 100644
index 0000000000000000000000000000000000000000..c90240b6c8173c3852c69797f319e3908d0dd1d4
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/658f2dccd9887c974b4f574fe0511442 differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/77fa6ec81fda02b6d41760cf6d403cf1 b/checkpoint_400000/ocdbt.process_0/d/77fa6ec81fda02b6d41760cf6d403cf1
new file mode 100644
index 0000000000000000000000000000000000000000..9bb9ccd3b5bf8bd793419034ec420cf11298b7b4
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/77fa6ec81fda02b6d41760cf6d403cf1 differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/9516c40bd7a380adac5cb4ca3997c11a b/checkpoint_400000/ocdbt.process_0/d/9516c40bd7a380adac5cb4ca3997c11a
new file mode 100644
index 0000000000000000000000000000000000000000..5dc02caebe0a67990492a6a745eb868c4e1c7ce2
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/9516c40bd7a380adac5cb4ca3997c11a differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/b137f1e9ebe14cb480bc3fa3fcfd6e1a b/checkpoint_400000/ocdbt.process_0/d/b137f1e9ebe14cb480bc3fa3fcfd6e1a
new file mode 100644
index 0000000000000000000000000000000000000000..05319277551aa602b75554a46df55b6d5351ba9f
--- /dev/null
+++ b/checkpoint_400000/ocdbt.process_0/d/b137f1e9ebe14cb480bc3fa3fcfd6e1a
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f3e836477d98666a3c1b2feda9d6b0fba58d276f521ff56aa821bdcf81db562
+size 228667
diff --git a/checkpoint_400000/ocdbt.process_0/d/b4ced3ec77dd597eb933791f7b761a94 b/checkpoint_400000/ocdbt.process_0/d/b4ced3ec77dd597eb933791f7b761a94
new file mode 100644
index 0000000000000000000000000000000000000000..12495005acb06fc0ed59b307ae462cea0ba27f40
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/b4ced3ec77dd597eb933791f7b761a94 differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/e8eb9bccc04d635a1af744e54c2ecbc8 b/checkpoint_400000/ocdbt.process_0/d/e8eb9bccc04d635a1af744e54c2ecbc8
new file mode 100644
index 0000000000000000000000000000000000000000..0f0ed433e05f78a31db016fc9a7f8b9a157d126d
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/e8eb9bccc04d635a1af744e54c2ecbc8 differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/f4c5a668a18b620ff8fee12ab7036abe b/checkpoint_400000/ocdbt.process_0/d/f4c5a668a18b620ff8fee12ab7036abe
new file mode 100644
index 0000000000000000000000000000000000000000..738131740f79d1d3668e478ce53aad3f2c4141c0
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/f4c5a668a18b620ff8fee12ab7036abe differ
diff --git a/checkpoint_400000/ocdbt.process_0/d/fe42395121a7437ea74c9a147db5570d b/checkpoint_400000/ocdbt.process_0/d/fe42395121a7437ea74c9a147db5570d
new file mode 100644
index 0000000000000000000000000000000000000000..bca5db1da223105988efe317cbae7458a9fc9b22
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/d/fe42395121a7437ea74c9a147db5570d differ
diff --git a/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..deee4296d9e77b7d076ed614ad819d0effc68f6d
Binary files /dev/null and b/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/latest/checkpoint_400000/_CHECKPOINT_METADATA b/latest/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..9e66302113bf504077a9654111eb76471cfd47a8
--- /dev/null
+++ b/latest/checkpoint_400000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778652340133491234, "commit_timestamp_nsecs": 1778652356572359527, "custom_metadata": {}}
\ No newline at end of file
diff --git a/latest/checkpoint_400000/_METADATA b/latest/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..bbd0485d06d3a08df040be11c549f90fb114ae53
--- /dev/null
+++ b/latest/checkpoint_400000/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('backbone', 'DiTBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_10', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_10', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_11', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_11', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_12', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_12', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_12', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_12', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_13', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_13', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_13', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_13', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_14', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_14', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_14', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_14', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_15', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_15', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_15', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_15', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_16', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_16', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_16', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_16', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_17', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_17', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_17', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_17', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_18', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_18', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_18', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_18', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_19', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_19', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_19', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_19', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_2', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_2', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_20', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_20', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_20', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_20', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_21', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_21', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_21', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_21', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_22', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_22', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_22', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_22', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_22', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_22', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_23', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_23', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_23', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_23', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_23', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_23', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_24', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_24', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_24', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_24', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_24', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_24', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_25', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_25', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_25', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_25', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_25', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_25', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_26', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_26', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_26', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_26', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_26', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_26', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_27', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_27', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_27', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_27', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_27', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_27', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_3', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_3', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_4', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_4', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_5', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_5', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_6', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_6', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_7', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_7', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_8', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_8', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('backbone', 'DiTBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('backbone', 'DiTBlock_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('backbone', 'DiTBlock_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('backbone', 'DiTBlock_9', 'Dense_2', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_9', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('backbone', 'FinalLayer_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('backbone', 'FinalLayer_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('backbone', 'FinalLayer_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304]}}, "('backbone', 'FinalLayer_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 2304]}}, "('backbone', 'LabelEmbedder_0', 'Embed_0', 'embedding')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "LabelEmbedder_0", "key_type": 2}, {"key": "Embed_0", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('backbone', 'PatchedPatchEmbed_0', 'proj', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'PatchedPatchEmbed_0', 'proj', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 1152]}}, "('backbone', 'TimestepEmbedder_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'TimestepEmbedder_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1152]}}, "('backbone', 'TimestepEmbedder_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('backbone', 'TimestepEmbedder_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 1152]}}, "('predictor', 'blocks_0', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_0', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_0', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_0', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_0', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_0', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_0', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_0', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_0', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_0', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_0', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_0', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_1', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_1', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_1', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_1', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_1', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_1', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_1', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_1', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_1', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_1', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_1', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_1', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_10', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_10', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_10', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_10', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_10', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_10', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_10', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_10', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_10', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_10', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_10', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_10', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_11', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_11', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_11', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_11', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('predictor', 'blocks_11', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_11', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_11', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_11', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_11', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_11', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_11', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_11', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_11', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_11', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_11', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_11', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_11', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_11', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_11', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_11', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_11', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_11', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_2', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_2', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_2', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_2', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_2', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_2', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_2', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_2', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_2', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_2', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_2', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_2', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_3', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_3', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_3', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_3', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('predictor', 'blocks_3', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_3', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_3', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_3', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_3', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_3', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_3', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_3', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_3', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_3', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_3', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_3', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_3', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_3', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_3', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_3', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_3', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_3', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_4', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_4', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_4', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_4', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_4', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_4', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_4', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_4', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_4', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_4', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_4', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_4', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_5', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_5', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_5', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_5', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_5', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_5', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_5', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_5', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_5', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_5', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_5', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_5', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_6', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_6', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_6', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_6', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_6', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_6', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_6', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_6', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_6', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_6', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_6', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_6', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_7', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_7', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_7', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_7', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('predictor', 'blocks_7', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_7', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_7', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('predictor', 'blocks_7', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('predictor', 'blocks_7', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_7', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_7', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_7', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_7', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_7', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_7', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_7', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_7', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_7', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_7', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_7', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_7', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_7', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_8', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_8', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_8', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_8', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_8', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_8', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_8', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_8', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_8', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_8', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_8', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_8', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'blocks_9', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_9', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_9', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_9', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('predictor', 'blocks_9', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('predictor', 'blocks_9', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('predictor', 'blocks_9', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('predictor', 'blocks_9', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('predictor', 'blocks_9', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_9', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('predictor', 'blocks_9', 'pwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'blocks_9', 'pwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('predictor', 'class_input_embed', 'embedding')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "class_input_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('predictor', 'cond_delta_embed', 'embedding')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_delta_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [28, 32]}}, "('predictor', 'cond_layer_embed', 'embedding')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_layer_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [29, 32]}}, "('predictor', 'cond_out', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('predictor', 'cond_out', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('predictor', 'cond_t_proj', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('predictor', 'cond_t_proj', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('predictor', 'final_ln', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'final_ln', 'scale')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'gamma_out')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "gamma_out", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('predictor', 'in_proj', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('predictor', 'in_proj', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 480]}}, "('predictor', 'mag_head', 'adaln_mod', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [964]}}, "('predictor', 'mag_head', 'adaln_mod', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 964]}}, "('predictor', 'mag_head', 'dwconv', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482]}}, "('predictor', 'mag_head', 'dwconv', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 482]}}, "('predictor', 'mag_head', 'pw1', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120]}}, "('predictor', 'mag_head', 'pw1', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482, 120]}}, "('predictor', 'mag_head', 'pw2', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('predictor', 'mag_head', 'pw2', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120, 1]}}, "('predictor', 'out_proj', 'bias')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('predictor', 'out_proj', 'kernel')": {"key_metadata": [{"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1152]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/latest/checkpoint_400000/_sharding b/latest/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..384403235a88279737437bfee68391867a2b8108
--- /dev/null
+++ b/latest/checkpoint_400000/_sharding
@@ -0,0 +1 @@
+{"YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRmluYWxMYXllcl8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRmluYWxMYXllcl8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRmluYWxMYXllcl8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuRmluYWxMYXllcl8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuTGFiZWxFbWJlZGRlcl8wLkVtYmVkXzAuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuUGF0Y2hlZFBhdGNoRW1iZWRfMC5wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuUGF0Y2hlZFBhdGNoRW1iZWRfMC5wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuVGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuVGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuVGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmFja2JvbmUuVGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLmFkYWxuX21vZC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLmFkYWxuX21vZC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLnB3MS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLnB3MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLnB3Mi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm1hZ19oZWFkLnB3Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm91dF9wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLm91dF9wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc180LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc181LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc182LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc183LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc184LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc185LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18wLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5jb252X2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5jb252X2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5kd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5kd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5wd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMC5wd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5jb252X2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5jb252X2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5kd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5kd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5wd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18xMS5wd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18yLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmJsb2Nrc18zLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNsYXNzX2lucHV0X2VtYmVkLmVtYmVkZGluZw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfZGVsdGFfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfb3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfb3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfbGF5ZXJfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfdF9wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmNvbmRfdF9wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmZpbmFsX2xuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmZpbmFsX2xuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmdhbW1hX291dA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmluX3Byb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","cHJlZGljdG9yLmluX3Byb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}"}
\ No newline at end of file
diff --git a/latest/checkpoint_400000/array_metadatas/process_0 b/latest/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..d267d96a4c3b8ebf13b6f6b35fac282dadea1079
--- /dev/null
+++ b/latest/checkpoint_400000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_24.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_25.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_26.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_27.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_0.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_1.bias", "write_shape": [2304], "chunk_shape": [2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_1.kernel", "write_shape": [1152, 2304], "chunk_shape": [1152, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.LabelEmbedder_0.Embed_0.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.PatchedPatchEmbed_0.proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.PatchedPatchEmbed_0.proj.kernel", "write_shape": [16, 1152], "chunk_shape": [16, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_0.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_0.kernel", "write_shape": [256, 1152], "chunk_shape": [256, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_1.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_1.kernel", "write_shape": [1152, 1152], "chunk_shape": [1152, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_0.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_1.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_10.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_11.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_2.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_3.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_4.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_5.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_6.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_7.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_8.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.blocks_9.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.class_input_embed.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_delta_embed.embedding", "write_shape": [28, 32], "chunk_shape": [28, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_layer_embed.embedding", "write_shape": [29, 32], "chunk_shape": [29, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_out.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_out.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_t_proj.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_t_proj.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.final_ln.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.final_ln.scale", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.gamma_out", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.in_proj.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.in_proj.kernel", "write_shape": [1152, 480], "chunk_shape": [1152, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.adaln_mod.bias", "write_shape": [964], "chunk_shape": [964], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.adaln_mod.kernel", "write_shape": [32, 964], "chunk_shape": [32, 964], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.dwconv.bias", "write_shape": [482], "chunk_shape": [482], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.dwconv.kernel", "write_shape": [3, 3, 1, 482], "chunk_shape": [3, 3, 1, 482], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.pw1.bias", "write_shape": [120], "chunk_shape": [120], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.pw1.kernel", "write_shape": [482, 120], "chunk_shape": [482, 120], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.pw2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.mag_head.pw2.kernel", "write_shape": [120, 1], "chunk_shape": [120, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.out_proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.out_proj.kernel", "write_shape": [480, 1152], "chunk_shape": [480, 1152], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/latest/checkpoint_400000/d/7dce239a0396b9174b4e202f8dc8aa6d b/latest/checkpoint_400000/d/7dce239a0396b9174b4e202f8dc8aa6d
new file mode 100644
index 0000000000000000000000000000000000000000..67853fb391bb62606381293f7028292c24e83198
Binary files /dev/null and b/latest/checkpoint_400000/d/7dce239a0396b9174b4e202f8dc8aa6d differ
diff --git a/latest/checkpoint_400000/manifest.ocdbt b/latest/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..014ccc382711b7441f1a668bd9c098e1b378c8e8
Binary files /dev/null and b/latest/checkpoint_400000/manifest.ocdbt differ
diff --git a/latest/checkpoint_400000/ocdbt.process_0/d/48ae8c3b25732f35127dee2411fe8c0b b/latest/checkpoint_400000/ocdbt.process_0/d/48ae8c3b25732f35127dee2411fe8c0b
new file mode 100644
index 0000000000000000000000000000000000000000..b53a9d4930c60291bac09507b9cedfdbbd88e123
Binary files /dev/null and b/latest/checkpoint_400000/ocdbt.process_0/d/48ae8c3b25732f35127dee2411fe8c0b differ
diff --git a/latest/checkpoint_400000/ocdbt.process_0/d/62d5cced42ccdbcfef8aa675459c669c b/latest/checkpoint_400000/ocdbt.process_0/d/62d5cced42ccdbcfef8aa675459c669c
new file mode 100644
index 0000000000000000000000000000000000000000..2ce5c9f62d286c50a601e32c12445ddd422d2956
Binary files /dev/null and b/latest/checkpoint_400000/ocdbt.process_0/d/62d5cced42ccdbcfef8aa675459c669c differ
diff --git a/latest/checkpoint_400000/ocdbt.process_0/d/d6c443f7c7bdb92c327f60c696f6762c b/latest/checkpoint_400000/ocdbt.process_0/d/d6c443f7c7bdb92c327f60c696f6762c
new file mode 100644
index 0000000000000000000000000000000000000000..9ccbde51cf9e24f8fb8cfdce408a136378ba4262
Binary files /dev/null and b/latest/checkpoint_400000/ocdbt.process_0/d/d6c443f7c7bdb92c327f60c696f6762c differ
diff --git a/latest/checkpoint_400000/ocdbt.process_0/d/f9a15bfe79964f926d494f62c37a2d11 b/latest/checkpoint_400000/ocdbt.process_0/d/f9a15bfe79964f926d494f62c37a2d11
new file mode 100644
index 0000000000000000000000000000000000000000..0cd5ec055383513b542e322541ccb4766f155956
Binary files /dev/null and b/latest/checkpoint_400000/ocdbt.process_0/d/f9a15bfe79964f926d494f62c37a2d11 differ
diff --git a/latest/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/latest/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..facbd29a3c7be06c419b408143825bb5261876e1
Binary files /dev/null and b/latest/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/latest/ema/checkpoint_400000/_CHECKPOINT_METADATA b/latest/ema/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..355232fbbd03d7ed8b75aec597fc3eb4020871dd
--- /dev/null
+++ b/latest/ema/checkpoint_400000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778652390965831684, "commit_timestamp_nsecs": 1778652407047983906, "custom_metadata": {}}
\ No newline at end of file
diff --git a/latest/ema/checkpoint_400000/_METADATA b/latest/ema/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..5584b786c2734b34d4f42ff3c34bdda85736e516
--- /dev/null
+++ b/latest/ema/checkpoint_400000/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('DiTBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_10', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_10', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_11', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_11', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_12', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_12', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_12', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_12', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_13', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_13', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_13', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_13', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_14', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_14', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_14', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_14', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_15', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_15', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_15', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_15', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_16', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_16', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_16', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_16', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_17', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_17', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_17', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_17', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_18', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_18', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_18', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_18', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_19', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_19', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_19', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_19', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_2', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_2', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_20', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_20', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_20', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_20', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_21', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_21', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_21', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_21', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_22', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_22', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_22', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_22', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_22', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_22', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_23', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_23', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_23', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_23', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_23', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_23', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_24', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_24', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_24', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_24', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_24', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_24', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_25', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_25', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_25', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_25', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_25', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_25', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_26', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_26', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_26', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_26', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_26', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_26', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_27', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_27', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_27', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_27', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_27', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_27', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_3', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_3', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_4', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_4', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_5', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_5', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_6', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_6', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_7', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_7', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_8', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_8', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('DiTBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('DiTBlock_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('DiTBlock_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('DiTBlock_9', 'Dense_2', 'bias')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_9', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('FinalLayer_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('FinalLayer_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('FinalLayer_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304]}}, "('FinalLayer_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 2304]}}, "('LabelEmbedder_0', 'Embed_0', 'embedding')": {"key_metadata": [{"key": "LabelEmbedder_0", "key_type": 2}, {"key": "Embed_0", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('PatchedPatchEmbed_0', 'proj', 'bias')": {"key_metadata": [{"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('PatchedPatchEmbed_0', 'proj', 'kernel')": {"key_metadata": [{"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 1152]}}, "('TimestepEmbedder_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('TimestepEmbedder_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1152]}}, "('TimestepEmbedder_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('TimestepEmbedder_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 1152]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/latest/ema/checkpoint_400000/_sharding b/latest/ema/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..c85ea519376bbc045b3b6ce4ae03afd7a261c561
--- /dev/null
+++ b/latest/ema/checkpoint_400000/_sharding
@@ -0,0 +1 @@
+{"RGlUQmxvY2tfMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTguTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMTkuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjAuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjEuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjIuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjMuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjQuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjUuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjYuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuRGVuc2VfMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuRGVuc2VfMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMjcuTXVsdGlIZWFkRG90UHJvZHVjdEF0dGVudGlvbl8wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RGlUQmxvY2tfOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RmluYWxMYXllcl8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RmluYWxMYXllcl8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RmluYWxMYXllcl8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","RmluYWxMYXllcl8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","TGFiZWxFbWJlZGRlcl8wLkVtYmVkXzAuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","UGF0Y2hlZFBhdGNoRW1iZWRfMC5wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","UGF0Y2hlZFBhdGNoRW1iZWRfMC5wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","VGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","VGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","VGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","VGltZXN0ZXBFbWJlZGRlcl8wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}"}
\ No newline at end of file
diff --git a/latest/ema/checkpoint_400000/array_metadatas/process_0 b/latest/ema/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..b2b062dde4cf1b53382d65e8bdfe237daa27d161
--- /dev/null
+++ b/latest/ema/checkpoint_400000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "DiTBlock_0.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_24.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_25.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_26.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_27.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "FinalLayer_0.Dense_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "FinalLayer_0.Dense_0.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "FinalLayer_0.Dense_1.bias", "write_shape": [2304], "chunk_shape": [2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "FinalLayer_0.Dense_1.kernel", "write_shape": [1152, 2304], "chunk_shape": [1152, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "LabelEmbedder_0.Embed_0.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "PatchedPatchEmbed_0.proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "PatchedPatchEmbed_0.proj.kernel", "write_shape": [16, 1152], "chunk_shape": [16, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "TimestepEmbedder_0.Dense_0.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "TimestepEmbedder_0.Dense_0.kernel", "write_shape": [256, 1152], "chunk_shape": [256, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "TimestepEmbedder_0.Dense_1.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "TimestepEmbedder_0.Dense_1.kernel", "write_shape": [1152, 1152], "chunk_shape": [1152, 1152], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/latest/ema/checkpoint_400000/d/9989ed5404cf40f857105d060af10e5c b/latest/ema/checkpoint_400000/d/9989ed5404cf40f857105d060af10e5c
new file mode 100644
index 0000000000000000000000000000000000000000..8148f96917388f2b6a242eb0028997e8bc5d560b
Binary files /dev/null and b/latest/ema/checkpoint_400000/d/9989ed5404cf40f857105d060af10e5c differ
diff --git a/latest/ema/checkpoint_400000/manifest.ocdbt b/latest/ema/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..bff3ffdc6bbcf906d0cb57cf473358b0c7c9e29b
Binary files /dev/null and b/latest/ema/checkpoint_400000/manifest.ocdbt differ
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/d/0b6be26a286db035b0dced5cc40d0217 b/latest/ema/checkpoint_400000/ocdbt.process_0/d/0b6be26a286db035b0dced5cc40d0217
new file mode 100644
index 0000000000000000000000000000000000000000..e85b1791230264243f94cd93800e1a5818887eef
Binary files /dev/null and b/latest/ema/checkpoint_400000/ocdbt.process_0/d/0b6be26a286db035b0dced5cc40d0217 differ
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/d/1ab113fa74cd4906c1f19a46c03c4314 b/latest/ema/checkpoint_400000/ocdbt.process_0/d/1ab113fa74cd4906c1f19a46c03c4314
new file mode 100644
index 0000000000000000000000000000000000000000..55e357d6f288bbd41b328606c095437fa711d650
Binary files /dev/null and b/latest/ema/checkpoint_400000/ocdbt.process_0/d/1ab113fa74cd4906c1f19a46c03c4314 differ
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/d/1de9dc33f844966b10fa8a34bb366867 b/latest/ema/checkpoint_400000/ocdbt.process_0/d/1de9dc33f844966b10fa8a34bb366867
new file mode 100644
index 0000000000000000000000000000000000000000..e4922e416dbd003330f3f35210389ea7c8d00047
--- /dev/null
+++ b/latest/ema/checkpoint_400000/ocdbt.process_0/d/1de9dc33f844966b10fa8a34bb366867
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:154ace00278dcc783c63813a61720aff44f9561f6a94e3560c721525e77b0036
+size 59273216
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/d/21156d2763b1a84f25cc1b9c9b64ec35 b/latest/ema/checkpoint_400000/ocdbt.process_0/d/21156d2763b1a84f25cc1b9c9b64ec35
new file mode 100644
index 0000000000000000000000000000000000000000..a705c155a53292920e35c481e5a2fefb884c2032
Binary files /dev/null and b/latest/ema/checkpoint_400000/ocdbt.process_0/d/21156d2763b1a84f25cc1b9c9b64ec35 differ
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/d/4206e9dce6a03899b0e39cfd3720d1da b/latest/ema/checkpoint_400000/ocdbt.process_0/d/4206e9dce6a03899b0e39cfd3720d1da
new file mode 100644
index 0000000000000000000000000000000000000000..f6c8a18cad6a1b7b614879bf7d2f8ef0c8c33a9c
Binary files /dev/null and b/latest/ema/checkpoint_400000/ocdbt.process_0/d/4206e9dce6a03899b0e39cfd3720d1da differ
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/d/44935abe029c8ac7e52c7a04662b78f8 b/latest/ema/checkpoint_400000/ocdbt.process_0/d/44935abe029c8ac7e52c7a04662b78f8
new file mode 100644
index 0000000000000000000000000000000000000000..f2814986fd131ca22e3bd8fe3b153a12a4c38ce5
Binary files /dev/null and b/latest/ema/checkpoint_400000/ocdbt.process_0/d/44935abe029c8ac7e52c7a04662b78f8 differ
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/d/6dd5db8b66e279b58821d30fb0d98362 b/latest/ema/checkpoint_400000/ocdbt.process_0/d/6dd5db8b66e279b58821d30fb0d98362
new file mode 100644
index 0000000000000000000000000000000000000000..9572b854cb7dbaeb4f739f5b460c50691ec1c9c0
Binary files /dev/null and b/latest/ema/checkpoint_400000/ocdbt.process_0/d/6dd5db8b66e279b58821d30fb0d98362 differ
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/d/80d4d7186e25094c2a630a6d91f24a8f b/latest/ema/checkpoint_400000/ocdbt.process_0/d/80d4d7186e25094c2a630a6d91f24a8f
new file mode 100644
index 0000000000000000000000000000000000000000..447cd0b282931ea34354a3135be8f84678507278
--- /dev/null
+++ b/latest/ema/checkpoint_400000/ocdbt.process_0/d/80d4d7186e25094c2a630a6d91f24a8f
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ffcdf7aa6e85199557356255c65a2a5c25392f2f6cc4f085878a0bf223dadc82
+size 4960256
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/d/a1c1bbbffe494f5fb8bdbb4355efeeb6 b/latest/ema/checkpoint_400000/ocdbt.process_0/d/a1c1bbbffe494f5fb8bdbb4355efeeb6
new file mode 100644
index 0000000000000000000000000000000000000000..5b82f6138e9e2f3985c48b8acaaca02782fec050
Binary files /dev/null and b/latest/ema/checkpoint_400000/ocdbt.process_0/d/a1c1bbbffe494f5fb8bdbb4355efeeb6 differ
diff --git a/latest/ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/latest/ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..7cfe75c813902ab573a30e00ea113f7dacbf4451
Binary files /dev/null and b/latest/ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/latest/l2_ema/checkpoint_400000/_CHECKPOINT_METADATA b/latest/l2_ema/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..cbd5670419f21a4354ce5e72a27c2a920e116545
--- /dev/null
+++ b/latest/l2_ema/checkpoint_400000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778652408080559757, "commit_timestamp_nsecs": 1778652408101433400, "custom_metadata": {}}
\ No newline at end of file
diff --git a/latest/l2_ema/checkpoint_400000/_METADATA b/latest/l2_ema/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..ada3cd782e85f12ef9574e31944245c2e7aefd59
--- /dev/null
+++ b/latest/l2_ema/checkpoint_400000/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('l2_ema',)": {"key_metadata": [{"key": "l2_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [29, 5, 256]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/latest/l2_ema/checkpoint_400000/_sharding b/latest/l2_ema/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..ac31b218b77c613e621eacb64ab85a4c534b1351
--- /dev/null
+++ b/latest/l2_ema/checkpoint_400000/_sharding
@@ -0,0 +1 @@
+{"bDJfZW1h":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}"}
\ No newline at end of file
diff --git a/latest/l2_ema/checkpoint_400000/array_metadatas/process_0 b/latest/l2_ema/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..e9de3c50745e6b0030f4c1a5806ad64586a79191
--- /dev/null
+++ b/latest/l2_ema/checkpoint_400000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "l2_ema", "write_shape": [29, 5, 256], "chunk_shape": [29, 5, 256], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/latest/l2_ema/checkpoint_400000/d/29a7bf31379c65d6c3abc180d99c5f16 b/latest/l2_ema/checkpoint_400000/d/29a7bf31379c65d6c3abc180d99c5f16
new file mode 100644
index 0000000000000000000000000000000000000000..6814fdd146ab082b68af098da9b00d429c6fad81
Binary files /dev/null and b/latest/l2_ema/checkpoint_400000/d/29a7bf31379c65d6c3abc180d99c5f16 differ
diff --git a/latest/l2_ema/checkpoint_400000/manifest.ocdbt b/latest/l2_ema/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..520072e5dd7f9c91bc99cf95ddd0f3c9bfc62fbf
Binary files /dev/null and b/latest/l2_ema/checkpoint_400000/manifest.ocdbt differ
diff --git a/latest/l2_ema/checkpoint_400000/ocdbt.process_0/d/0f1593ea9937707c8cc76cc134300dcf b/latest/l2_ema/checkpoint_400000/ocdbt.process_0/d/0f1593ea9937707c8cc76cc134300dcf
new file mode 100644
index 0000000000000000000000000000000000000000..48b58b703c5ad4db8875b7f232584b0d21c71069
Binary files /dev/null and b/latest/l2_ema/checkpoint_400000/ocdbt.process_0/d/0f1593ea9937707c8cc76cc134300dcf differ
diff --git a/latest/l2_ema/checkpoint_400000/ocdbt.process_0/d/ad0a27e00df468000d831d7c62d313d1 b/latest/l2_ema/checkpoint_400000/ocdbt.process_0/d/ad0a27e00df468000d831d7c62d313d1
new file mode 100644
index 0000000000000000000000000000000000000000..6814fdd146ab082b68af098da9b00d429c6fad81
Binary files /dev/null and b/latest/l2_ema/checkpoint_400000/ocdbt.process_0/d/ad0a27e00df468000d831d7c62d313d1 differ
diff --git a/latest/l2_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/latest/l2_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..725f3b0a0ebf0f864217f1c29c2672a2511ee034
Binary files /dev/null and b/latest/l2_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/latest/opt_state/checkpoint_400000/_CHECKPOINT_METADATA b/latest/opt_state/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..c90b9759be17d85a158997fed5afc846f465a09e
--- /dev/null
+++ b/latest/opt_state/checkpoint_400000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778652356903073194, "commit_timestamp_nsecs": 1778652390334628953, "custom_metadata": {}}
\ No newline at end of file
diff --git a/latest/opt_state/checkpoint_400000/_METADATA b/latest/opt_state/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..886fdce288e71cd7f56d1b4f295a74f92132bce4
--- /dev/null
+++ b/latest/opt_state/checkpoint_400000/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('inner_states', 'backbone', 'inner_state', '0')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'count')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'FinalLayer_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'FinalLayer_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'FinalLayer_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'FinalLayer_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 2304]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'LabelEmbedder_0', 'Embed_0', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "LabelEmbedder_0", "key_type": 2}, {"key": "Embed_0", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'PatchedPatchEmbed_0', 'proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'PatchedPatchEmbed_0', 'proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'TimestepEmbedder_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'TimestepEmbedder_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'TimestepEmbedder_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'backbone', 'TimestepEmbedder_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'class_input_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "class_input_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_delta_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_delta_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_layer_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_layer_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_t_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_t_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'final_ln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'final_ln', 'scale')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'gamma_out')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "gamma_out", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'in_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'in_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'adaln_mod', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'adaln_mod', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'pw1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'pw1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'pw2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'pw2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'out_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'mu', 'predictor', 'out_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 6912]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 4608]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4608, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 16, 72]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'FinalLayer_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'FinalLayer_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'FinalLayer_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'FinalLayer_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 2304]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'LabelEmbedder_0', 'Embed_0', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "LabelEmbedder_0", "key_type": 2}, {"key": "Embed_0", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'PatchedPatchEmbed_0', 'proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'PatchedPatchEmbed_0', 'proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'TimestepEmbedder_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'TimestepEmbedder_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'TimestepEmbedder_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'backbone', 'TimestepEmbedder_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 1152]}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'class_input_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "class_input_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_delta_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_delta_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_layer_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_layer_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_t_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_t_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'final_ln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'final_ln', 'scale')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'gamma_out')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "gamma_out", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'in_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'in_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'adaln_mod', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'adaln_mod', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'pw1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'pw1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'pw2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'pw2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'out_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '0', 'nu', 'predictor', 'out_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '1', 'inner_state')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "1", "key_type": 1}, {"key": "inner_state", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'backbone', 'inner_state', '1', '2')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "2", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '0')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'count')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'FinalLayer_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'FinalLayer_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'FinalLayer_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'FinalLayer_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'LabelEmbedder_0', 'Embed_0', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "LabelEmbedder_0", "key_type": 2}, {"key": "Embed_0", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'PatchedPatchEmbed_0', 'proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'PatchedPatchEmbed_0', 'proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'TimestepEmbedder_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'TimestepEmbedder_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'TimestepEmbedder_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'backbone', 'TimestepEmbedder_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_0', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_1', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_10', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_11', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_2', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_3', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_4', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_5', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_6', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_7', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_8', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'blocks_9', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'class_input_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "class_input_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_delta_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_delta_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [28, 32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_layer_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_layer_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [29, 32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_t_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'cond_t_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'final_ln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'final_ln', 'scale')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'gamma_out')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "gamma_out", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'in_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'in_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'adaln_mod', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [964]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'adaln_mod', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 964]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 482]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'pw1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'pw1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482, 120]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'pw2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'mag_head', 'pw2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120, 1]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'out_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'mu', 'predictor', 'out_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1152]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_0', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_0", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_1', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_1", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_10', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_10", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_11', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_11", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_12', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_12", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_13', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_13", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_14', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_14", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_15', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_15", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_16', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_16", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_17', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_17", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_18', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_18", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_19', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_19", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_2', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_2", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_20', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_20", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_21', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_21", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_22', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_22", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_23', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_23", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_24', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_24", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_25', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_25", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_26', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_26", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_27', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_27", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_3', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_3", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_4', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_4", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_5', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_5", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_6', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_6", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_7', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_7", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_8', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_8", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'DiTBlock_9', 'MultiHeadDotProductAttention_0', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "DiTBlock_9", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'FinalLayer_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'FinalLayer_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'FinalLayer_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'FinalLayer_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "FinalLayer_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'LabelEmbedder_0', 'Embed_0', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "LabelEmbedder_0", "key_type": 2}, {"key": "Embed_0", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'PatchedPatchEmbed_0', 'proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'PatchedPatchEmbed_0', 'proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "PatchedPatchEmbed_0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'TimestepEmbedder_0', 'Dense_0', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'TimestepEmbedder_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'TimestepEmbedder_0', 'Dense_1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'backbone', 'TimestepEmbedder_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "backbone", "key_type": 2}, {"key": "TimestepEmbedder_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_0', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_1', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_10', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_11', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_2', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_3', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_4', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_5', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_6', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_7', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_8', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'pwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'blocks_9', 'pwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'class_input_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "class_input_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_delta_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_delta_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [28, 32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_layer_embed', 'embedding')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_layer_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [29, 32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_out', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_out', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_t_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'cond_t_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "cond_t_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'final_ln', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'final_ln', 'scale')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "final_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'gamma_out')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "gamma_out", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'in_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'in_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 480]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'adaln_mod', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [964]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'adaln_mod', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 964]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'dwconv', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'dwconv', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 482]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'pw1', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'pw1', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482, 120]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'pw2', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'mag_head', 'pw2', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120, 1]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'out_proj', 'bias')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('inner_states', 'predictor', 'inner_state', '1', '0', 'nu', 'predictor', 'out_proj', 'kernel')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1152]}}, "('inner_states', 'predictor', 'inner_state', '1', '1')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('inner_states', 'predictor', 'inner_state', '1', '2')": {"key_metadata": [{"key": "inner_states", "key_type": 2}, {"key": "predictor", "key_type": 2}, {"key": "inner_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "2", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/latest/opt_state/checkpoint_400000/_sharding b/latest/opt_state/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..6e595a29f56e47685f073306931e3bc3c23fb29c
--- /dev/null
+++ b/latest/opt_state/checkpoint_400000/_sharding
@@ -0,0 +1 @@
+{"aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5jb3VudA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5GaW5hbExheWVyXzAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5GaW5hbExheWVyXzAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5GaW5hbExheWVyXzAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5GaW5hbExheWVyXzAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5MYWJlbEVtYmVkZGVyXzAuRW1iZWRfMC5lbWJlZGRpbmc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5QYXRjaGVkUGF0Y2hFbWJlZF8wLnByb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5QYXRjaGVkUGF0Y2hFbWJlZF8wLnByb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5UaW1lc3RlcEVtYmVkZGVyXzAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5UaW1lc3RlcEVtYmVkZGVyXzAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5UaW1lc3RlcEVtYmVkZGVyXzAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5tdS5iYWNrYm9uZS5UaW1lc3RlcEVtYmVkZGVyXzAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja180LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja181LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja182LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja183LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja184LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185Lk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja185LkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18wLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18xOS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yMy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNC5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNS5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNi5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8wLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8xLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8yLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18yNy5NdWx0aUhlYWREb3RQcm9kdWN0QXR0ZW50aW9uXzAudmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzAuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzAua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5EaVRCbG9ja18zLkRlbnNlXzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5GaW5hbExheWVyXzAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5GaW5hbExheWVyXzAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5GaW5hbExheWVyXzAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5GaW5hbExheWVyXzAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5MYWJlbEVtYmVkZGVyXzAuRW1iZWRfMC5lbWJlZGRpbmc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5QYXRjaGVkUGF0Y2hFbWJlZF8wLnByb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5QYXRjaGVkUGF0Y2hFbWJlZF8wLnByb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5UaW1lc3RlcEVtYmVkZGVyXzAuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5UaW1lc3RlcEVtYmVkZGVyXzAuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5UaW1lc3RlcEVtYmVkZGVyXzAuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLmJhY2tib25lLmlubmVyX3N0YXRlLjEuMC5udS5iYWNrYm9uZS5UaW1lc3RlcEVtYmVkZGVyXzAuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAuY291bnQ=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm1hZ19oZWFkLmFkYWxuX21vZC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm1hZ19oZWFkLmFkYWxuX21vZC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm1hZ19oZWFkLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm1hZ19oZWFkLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm1hZ19oZWFkLnB3MS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm1hZ19oZWFkLnB3MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm1hZ19oZWFkLnB3Mi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm1hZ19oZWFkLnB3Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm91dF9wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLm91dF9wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc180LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc181LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc182LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc183LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc184LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc185LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18wLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5jb252X2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5jb252X2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5kd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5kd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5wd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMC5wd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5jb252X2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5jb252X2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5kd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5kd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5wd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18xMS5wd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18yLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmJsb2Nrc18zLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmNsYXNzX2lucHV0X2VtYmVkLmVtYmVkZGluZw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmNvbmRfZGVsdGFfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmNvbmRfb3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmNvbmRfb3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmNvbmRfbGF5ZXJfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmNvbmRfdF9wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmNvbmRfdF9wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmZpbmFsX2xuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmZpbmFsX2xuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmdhbW1hX291dA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmluX3Byb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubXUucHJlZGljdG9yLmluX3Byb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm1hZ19oZWFkLmFkYWxuX21vZC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm1hZ19oZWFkLmFkYWxuX21vZC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm1hZ19oZWFkLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm1hZ19oZWFkLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm1hZ19oZWFkLnB3MS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm1hZ19oZWFkLnB3MS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm1hZ19oZWFkLnB3Mi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm1hZ19oZWFkLnB3Mi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm91dF9wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLm91dF9wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc180LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc181LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc182LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc183LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc184LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185Lm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185LmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185LmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185LmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185LmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185LnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc185LnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18wLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5jb252X2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5jb252X2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5kd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5kd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5tbHBfZmMyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5wd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMC5wd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLm91dC5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLm91dC5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLmtleS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLmtleS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5hdHRuX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5jb252X2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5jb252X2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5kd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5kd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5tbHBfZmMyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5wd2NvbnYuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18xMS5wd2NvbnYua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18yLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmJsb2Nrc18zLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmNsYXNzX2lucHV0X2VtYmVkLmVtYmVkZGluZw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmNvbmRfZGVsdGFfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmNvbmRfb3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmNvbmRfb3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmNvbmRfbGF5ZXJfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmNvbmRfdF9wcm9qLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmNvbmRfdF9wcm9qLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmZpbmFsX2xuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmZpbmFsX2xuLnNjYWxl":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmdhbW1hX291dA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmluX3Byb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5uZXJfc3RhdGVzLnByZWRpY3Rvci5pbm5lcl9zdGF0ZS4xLjAubnUucHJlZGljdG9yLmluX3Byb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}"}
\ No newline at end of file
diff --git a/latest/opt_state/checkpoint_400000/array_metadatas/process_0 b/latest/opt_state/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..872f85934cd8f81ef2417098b71a6886c60f3930
--- /dev/null
+++ b/latest/opt_state/checkpoint_400000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.FinalLayer_0.Dense_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.FinalLayer_0.Dense_0.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.FinalLayer_0.Dense_1.bias", "write_shape": [2304], "chunk_shape": [2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.FinalLayer_0.Dense_1.kernel", "write_shape": [1152, 2304], "chunk_shape": [1152, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.LabelEmbedder_0.Embed_0.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.PatchedPatchEmbed_0.proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.PatchedPatchEmbed_0.proj.kernel", "write_shape": [16, 1152], "chunk_shape": [16, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.TimestepEmbedder_0.Dense_0.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.TimestepEmbedder_0.Dense_0.kernel", "write_shape": [256, 1152], "chunk_shape": [256, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.TimestepEmbedder_0.Dense_1.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.mu.backbone.TimestepEmbedder_0.Dense_1.kernel", "write_shape": [1152, 1152], "chunk_shape": [1152, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_24.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_25.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_26.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_27.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.Dense_0.bias", "write_shape": [6912], "chunk_shape": [6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.Dense_0.kernel", "write_shape": [1152, 6912], "chunk_shape": [1152, 6912], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.Dense_1.bias", "write_shape": [4608], "chunk_shape": [4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.Dense_1.kernel", "write_shape": [1152, 4608], "chunk_shape": [1152, 4608], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.Dense_2.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.Dense_2.kernel", "write_shape": [4608, 1152], "chunk_shape": [4608, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 72, 1152], "chunk_shape": [16, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 72], "chunk_shape": [16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1152, 16, 72], "chunk_shape": [1152, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.FinalLayer_0.Dense_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.FinalLayer_0.Dense_0.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.FinalLayer_0.Dense_1.bias", "write_shape": [2304], "chunk_shape": [2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.FinalLayer_0.Dense_1.kernel", "write_shape": [1152, 2304], "chunk_shape": [1152, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.LabelEmbedder_0.Embed_0.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.PatchedPatchEmbed_0.proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.PatchedPatchEmbed_0.proj.kernel", "write_shape": [16, 1152], "chunk_shape": [16, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.TimestepEmbedder_0.Dense_0.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.TimestepEmbedder_0.Dense_0.kernel", "write_shape": [256, 1152], "chunk_shape": [256, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.TimestepEmbedder_0.Dense_1.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.backbone.inner_state.1.0.nu.backbone.TimestepEmbedder_0.Dense_1.kernel", "write_shape": [1152, 1152], "chunk_shape": [1152, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_0.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_1.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_10.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_11.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_2.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_3.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_4.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_5.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_6.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_7.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_8.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.blocks_9.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.class_input_embed.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.cond_delta_embed.embedding", "write_shape": [28, 32], "chunk_shape": [28, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.cond_layer_embed.embedding", "write_shape": [29, 32], "chunk_shape": [29, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.cond_out.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.cond_out.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.cond_t_proj.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.cond_t_proj.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.final_ln.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.final_ln.scale", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.gamma_out", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.in_proj.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.in_proj.kernel", "write_shape": [1152, 480], "chunk_shape": [1152, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.mag_head.adaln_mod.bias", "write_shape": [964], "chunk_shape": [964], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.mag_head.adaln_mod.kernel", "write_shape": [32, 964], "chunk_shape": [32, 964], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.mag_head.dwconv.bias", "write_shape": [482], "chunk_shape": [482], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.mag_head.dwconv.kernel", "write_shape": [3, 3, 1, 482], "chunk_shape": [3, 3, 1, 482], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.mag_head.pw1.bias", "write_shape": [120], "chunk_shape": [120], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.mag_head.pw1.kernel", "write_shape": [482, 120], "chunk_shape": [482, 120], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.mag_head.pw2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.mag_head.pw2.kernel", "write_shape": [120, 1], "chunk_shape": [120, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.out_proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.mu.predictor.out_proj.kernel", "write_shape": [480, 1152], "chunk_shape": [480, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_0.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_1.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_10.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_11.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_2.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_3.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_4.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_5.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_6.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_7.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_8.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.blocks_9.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.class_input_embed.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.cond_delta_embed.embedding", "write_shape": [28, 32], "chunk_shape": [28, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.cond_layer_embed.embedding", "write_shape": [29, 32], "chunk_shape": [29, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.cond_out.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.cond_out.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.cond_t_proj.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.cond_t_proj.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.final_ln.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.final_ln.scale", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.gamma_out", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.in_proj.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.in_proj.kernel", "write_shape": [1152, 480], "chunk_shape": [1152, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.mag_head.adaln_mod.bias", "write_shape": [964], "chunk_shape": [964], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.mag_head.adaln_mod.kernel", "write_shape": [32, 964], "chunk_shape": [32, 964], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.mag_head.dwconv.bias", "write_shape": [482], "chunk_shape": [482], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.mag_head.dwconv.kernel", "write_shape": [3, 3, 1, 482], "chunk_shape": [3, 3, 1, 482], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.mag_head.pw1.bias", "write_shape": [120], "chunk_shape": [120], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.mag_head.pw1.kernel", "write_shape": [482, 120], "chunk_shape": [482, 120], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.mag_head.pw2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.mag_head.pw2.kernel", "write_shape": [120, 1], "chunk_shape": [120, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.out_proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "inner_states.predictor.inner_state.1.0.nu.predictor.out_proj.kernel", "write_shape": [480, 1152], "chunk_shape": [480, 1152], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/latest/opt_state/checkpoint_400000/d/f3aa85c839f3b85d23b1d7b3459d4193 b/latest/opt_state/checkpoint_400000/d/f3aa85c839f3b85d23b1d7b3459d4193
new file mode 100644
index 0000000000000000000000000000000000000000..b74470013f393f4794efb3a771d29720184da385
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/d/f3aa85c839f3b85d23b1d7b3459d4193 differ
diff --git a/latest/opt_state/checkpoint_400000/manifest.ocdbt b/latest/opt_state/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..926c8b38db3952bedbace9e4d8c524b7cc95c908
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/manifest.ocdbt differ
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/02db637fa4bcc4215818a2a91e9d5932 b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/02db637fa4bcc4215818a2a91e9d5932
new file mode 100644
index 0000000000000000000000000000000000000000..a9706ffec8ab6c781527c2bc40464ce0bd389e79
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/02db637fa4bcc4215818a2a91e9d5932 differ
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/05a9473a63d559ee136609c88eae1f22 b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/05a9473a63d559ee136609c88eae1f22
new file mode 100644
index 0000000000000000000000000000000000000000..68f3cd38844597d078f177507d56d6cb0fb3b1e4
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/05a9473a63d559ee136609c88eae1f22 differ
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/3f10d27a3288c7011d5d10ef82202ad3 b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/3f10d27a3288c7011d5d10ef82202ad3
new file mode 100644
index 0000000000000000000000000000000000000000..f5ce1b84930c2952facbdaa633924c9402465468
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/3f10d27a3288c7011d5d10ef82202ad3 differ
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/55f24edf08ee69029b9ece4618cc887a b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/55f24edf08ee69029b9ece4618cc887a
new file mode 100644
index 0000000000000000000000000000000000000000..240cac3a013dcc6714307f24e4b7953474727d28
--- /dev/null
+++ b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/55f24edf08ee69029b9ece4618cc887a
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cd388aca4164328436e35fe8b395d66adaefe8ff8cb8dd4f68f4f88739a37ca
+size 5005312
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/66d489d7add771807d569474bb16bf69 b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/66d489d7add771807d569474bb16bf69
new file mode 100644
index 0000000000000000000000000000000000000000..1a0f5d7f998eeec26a7b1f3656f36f18eb55077f
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/66d489d7add771807d569474bb16bf69 differ
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/7800b6192bf894fc2a698b1c9a5bd3f9 b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/7800b6192bf894fc2a698b1c9a5bd3f9
new file mode 100644
index 0000000000000000000000000000000000000000..d83643c70fb4561c46c7c28569e3a06f3aff056b
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/7800b6192bf894fc2a698b1c9a5bd3f9 differ
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/90f764113abf7dadee2fb3110b0dd70f b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/90f764113abf7dadee2fb3110b0dd70f
new file mode 100644
index 0000000000000000000000000000000000000000..0da0a5ce8cfceeff607b794746411ef0338bef3b
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/90f764113abf7dadee2fb3110b0dd70f differ
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/a4424320b440d88fb60cb3ed03be89ce b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/a4424320b440d88fb60cb3ed03be89ce
new file mode 100644
index 0000000000000000000000000000000000000000..849203020b21d7027f24ccbec6ca03f850387f2b
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/a4424320b440d88fb60cb3ed03be89ce differ
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/fb8913f9a9fff237c42cd0f30f828b0e b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/fb8913f9a9fff237c42cd0f30f828b0e
new file mode 100644
index 0000000000000000000000000000000000000000..6a2a9715c10e1702586b95db29e8e05995f443f8
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/ocdbt.process_0/d/fb8913f9a9fff237c42cd0f30f828b0e differ
diff --git a/latest/opt_state/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/latest/opt_state/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..24f48d720ff2bdc2a6299932a6b42f06c1b5e85b
Binary files /dev/null and b/latest/opt_state/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/latest/predictor_ema/checkpoint_400000/_CHECKPOINT_METADATA b/latest/predictor_ema/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..b4c9554de32749d937656265b22cbb687ed6e604
--- /dev/null
+++ b/latest/predictor_ema/checkpoint_400000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778652407347446167, "commit_timestamp_nsecs": 1778652408056856560, "custom_metadata": {}}
\ No newline at end of file
diff --git a/latest/predictor_ema/checkpoint_400000/_METADATA b/latest/predictor_ema/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..98e40c42a0223d39dbd11e0b3d709b14663fc8bf
--- /dev/null
+++ b/latest/predictor_ema/checkpoint_400000/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('blocks_0', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_0', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_0', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_0', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_0', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_0', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_0', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_0', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_0', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_0', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_0', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_0', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_1', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_1', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_1', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_1', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_1', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_1', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_1', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_1', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_1', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_1', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_1', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_1', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_10', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_10', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_10', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_10', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_10', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_10', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_10', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_10', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_10', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_10', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_10', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_10', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_11', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_11', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_11', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_11', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('blocks_11', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_11', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_11', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_11', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_11', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_11', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_11', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_11', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_11', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_11', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_11', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_11', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_11', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_11', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_11', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_11', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_11', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_11', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_2', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_2', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_2', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_2', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_2', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_2', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_2', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_2', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_2', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_2', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_2', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_2', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_3', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_3', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_3', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_3', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('blocks_3', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_3', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_3', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_3', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_3', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_3', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_3', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_3', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_3', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_3', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_3', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_3', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_3', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_3', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_3', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_3', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_3', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_3', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_4', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_4', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_4', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_4', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_4', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_4', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_4', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_4', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_4', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_4', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_4', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_4', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_5', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_5', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_5', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_5', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_5', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_5', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_5', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_5', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_5', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_5', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_5', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_5', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_6', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_6', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_6', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_6', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_6', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_6', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_6', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_6', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_6', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_6', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_6', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_6', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_7', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_7', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_7', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_7', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('blocks_7', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_7', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_7', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_7', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_7', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_7', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_7', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_7', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_7', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_7', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_7', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_7', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_7', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_7', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_7', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_7', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_7', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_7', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_8', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_8', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_8', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_8', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_8', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_8', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_8', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_8', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_8', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_8', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_8', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_8', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_9', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_9', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_9', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_9', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_9', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_9', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_9', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_9', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_9', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_9', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_9', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_9', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('class_input_embed', 'embedding')": {"key_metadata": [{"key": "class_input_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('cond_delta_embed', 'embedding')": {"key_metadata": [{"key": "cond_delta_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [28, 32]}}, "('cond_layer_embed', 'embedding')": {"key_metadata": [{"key": "cond_layer_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [29, 32]}}, "('cond_out', 'bias')": {"key_metadata": [{"key": "cond_out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('cond_out', 'kernel')": {"key_metadata": [{"key": "cond_out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('cond_t_proj', 'bias')": {"key_metadata": [{"key": "cond_t_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('cond_t_proj', 'kernel')": {"key_metadata": [{"key": "cond_t_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('final_ln', 'bias')": {"key_metadata": [{"key": "final_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('final_ln', 'scale')": {"key_metadata": [{"key": "final_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('gamma_out',)": {"key_metadata": [{"key": "gamma_out", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('in_proj', 'bias')": {"key_metadata": [{"key": "in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('in_proj', 'kernel')": {"key_metadata": [{"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 480]}}, "('mag_head', 'adaln_mod', 'bias')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [964]}}, "('mag_head', 'adaln_mod', 'kernel')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 964]}}, "('mag_head', 'dwconv', 'bias')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482]}}, "('mag_head', 'dwconv', 'kernel')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 482]}}, "('mag_head', 'pw1', 'bias')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120]}}, "('mag_head', 'pw1', 'kernel')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482, 120]}}, "('mag_head', 'pw2', 'bias')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('mag_head', 'pw2', 'kernel')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120, 1]}}, "('out_proj', 'bias')": {"key_metadata": [{"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('out_proj', 'kernel')": {"key_metadata": [{"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1152]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/latest/predictor_ema/checkpoint_400000/_sharding b/latest/predictor_ema/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..e953c1dfb239a44e96eeebcd20a1eb878636ca2c
--- /dev/null
+++ b/latest/predictor_ema/checkpoint_400000/_sharding
@@ -0,0 +1 @@
+{"Y29uZF90X3Byb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF90X3Byb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF9kZWx0YV9lbWJlZC5lbWJlZGRpbmc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF9sYXllcl9lbWJlZC5lbWJlZGRpbmc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF9vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF9vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y2xhc3NfaW5wdXRfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzguY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzguY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzguZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzguZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Z2FtbWFfb3V0":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","ZmluYWxfbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","ZmluYWxfbG4uc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5fcHJvai5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5fcHJvai5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","b3V0X3Byb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","b3V0X3Byb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQuYWRhbG5fbW9kLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQuYWRhbG5fbW9kLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQucHcxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQucHcxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQucHcyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQucHcyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}"}
\ No newline at end of file
diff --git a/latest/predictor_ema/checkpoint_400000/array_metadatas/process_0 b/latest/predictor_ema/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..16aca7b61e88e945fe42c867d09c82f96f78d5d0
--- /dev/null
+++ b/latest/predictor_ema/checkpoint_400000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "blocks_0.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "class_input_embed.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_delta_embed.embedding", "write_shape": [28, 32], "chunk_shape": [28, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_layer_embed.embedding", "write_shape": [29, 32], "chunk_shape": [29, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_out.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_out.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_t_proj.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_t_proj.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "final_ln.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "final_ln.scale", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "gamma_out", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "in_proj.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "in_proj.kernel", "write_shape": [1152, 480], "chunk_shape": [1152, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.adaln_mod.bias", "write_shape": [964], "chunk_shape": [964], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.adaln_mod.kernel", "write_shape": [32, 964], "chunk_shape": [32, 964], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.dwconv.bias", "write_shape": [482], "chunk_shape": [482], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.dwconv.kernel", "write_shape": [3, 3, 1, 482], "chunk_shape": [3, 3, 1, 482], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.pw1.bias", "write_shape": [120], "chunk_shape": [120], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.pw1.kernel", "write_shape": [482, 120], "chunk_shape": [482, 120], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.pw2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.pw2.kernel", "write_shape": [120, 1], "chunk_shape": [120, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "out_proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "out_proj.kernel", "write_shape": [480, 1152], "chunk_shape": [480, 1152], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/latest/predictor_ema/checkpoint_400000/d/dbefcdd7bc0188d14c7084eb7fab19e0 b/latest/predictor_ema/checkpoint_400000/d/dbefcdd7bc0188d14c7084eb7fab19e0
new file mode 100644
index 0000000000000000000000000000000000000000..eb8bc6e39453ea52ebd675a7c523a504db06335f
Binary files /dev/null and b/latest/predictor_ema/checkpoint_400000/d/dbefcdd7bc0188d14c7084eb7fab19e0 differ
diff --git a/latest/predictor_ema/checkpoint_400000/manifest.ocdbt b/latest/predictor_ema/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..573a622105585deea574d20238f453c5ef17e9dd
Binary files /dev/null and b/latest/predictor_ema/checkpoint_400000/manifest.ocdbt differ
diff --git a/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/1e1dbbbd4fa63924d325bb36de9f1789 b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/1e1dbbbd4fa63924d325bb36de9f1789
new file mode 100644
index 0000000000000000000000000000000000000000..6f6a65bb2b23928b3ba3985682e0fea86e2d73b4
--- /dev/null
+++ b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/1e1dbbbd4fa63924d325bb36de9f1789
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9aac6bd38bedd6ebfce1fe63510c3168c74a5cec07f18b13c6b94fb0f6361f7
+size 102629376
diff --git a/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/721155ffa59b0423d52bc14d802e97f3 b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/721155ffa59b0423d52bc14d802e97f3
new file mode 100644
index 0000000000000000000000000000000000000000..849a6805ff0c202cd6aa4e7320ba28dab061ef8b
Binary files /dev/null and b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/721155ffa59b0423d52bc14d802e97f3 differ
diff --git a/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/9a268f7c8b2e0c77a42b67cd4920356f b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/9a268f7c8b2e0c77a42b67cd4920356f
new file mode 100644
index 0000000000000000000000000000000000000000..108594fda6360eee126ae566d6ab0eb0569930fa
--- /dev/null
+++ b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/9a268f7c8b2e0c77a42b67cd4920356f
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:49cd5537d66ff6aa97c76131c6dcd95a2145d1419698f348f1904842122ef41b
+size 179230
diff --git a/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/b5d754dc68b0a23b35350e348088431a b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/b5d754dc68b0a23b35350e348088431a
new file mode 100644
index 0000000000000000000000000000000000000000..7c91c0f5b3f0d61f67a37d041d19db984897a185
Binary files /dev/null and b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/b5d754dc68b0a23b35350e348088431a differ
diff --git a/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/c591d79ec4a696aa267327c4f7a5e49d b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/c591d79ec4a696aa267327c4f7a5e49d
new file mode 100644
index 0000000000000000000000000000000000000000..33a20b8f3335989adb12736e3124b264c590b4a1
Binary files /dev/null and b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/c591d79ec4a696aa267327c4f7a5e49d differ
diff --git a/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/ca22174a382c5e795a946cf230562125 b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/ca22174a382c5e795a946cf230562125
new file mode 100644
index 0000000000000000000000000000000000000000..84fac4e6538279867df449bfc02c49b9712219b0
--- /dev/null
+++ b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/ca22174a382c5e795a946cf230562125
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d55bb5184d595046fe7431f1f601468461916ced571e354187edb26c2cc4e4d
+size 13697024
diff --git a/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/daaeadea7f38c6e9f10d1dcc134f4505 b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/daaeadea7f38c6e9f10d1dcc134f4505
new file mode 100644
index 0000000000000000000000000000000000000000..ac5b72641d03483a0f104363384bd6cb4a12c9b8
Binary files /dev/null and b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/d/daaeadea7f38c6e9f10d1dcc134f4505 differ
diff --git a/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..0d3540814b8dbdb8126f6e23eaf5eff6fca90023
Binary files /dev/null and b/latest/predictor_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/opt_state/checkpoint_400000/ocdbt.process_0/d/4e81829b310436056edf3761802dc154 b/opt_state/checkpoint_400000/ocdbt.process_0/d/4e81829b310436056edf3761802dc154
new file mode 100644
index 0000000000000000000000000000000000000000..f3333ee47867714ab917ef7da49822aa48228046
--- /dev/null
+++ b/opt_state/checkpoint_400000/ocdbt.process_0/d/4e81829b310436056edf3761802dc154
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08a81820f4914e92ca2c2f92cd14c1ab6e28b61d615ed9b74e941b82e58c8762
+size 217829
diff --git a/permanent/step_100000/checkpoint_100000/ocdbt.process_0/d/ecf6faa66144dab60e1504644745b153 b/permanent/step_100000/checkpoint_100000/ocdbt.process_0/d/ecf6faa66144dab60e1504644745b153
new file mode 100644
index 0000000000000000000000000000000000000000..83ab1d8930ee3ad80cfebf5a1d55e0c8279f9753
--- /dev/null
+++ b/permanent/step_100000/checkpoint_100000/ocdbt.process_0/d/ecf6faa66144dab60e1504644745b153
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f1e56e5c36a0cfd463ca8693c05b229f6bf7b8a977b44991cb924308ba3f778
+size 363527
diff --git a/permanent/step_100000/ema/checkpoint_100000/ocdbt.process_0/d/5cdb2bacecd19679b6575ae33490dd36 b/permanent/step_100000/ema/checkpoint_100000/ocdbt.process_0/d/5cdb2bacecd19679b6575ae33490dd36
new file mode 100644
index 0000000000000000000000000000000000000000..d3085c5f79cb4ce24e54ca7eb60e1280af09b718
--- /dev/null
+++ b/permanent/step_100000/ema/checkpoint_100000/ocdbt.process_0/d/5cdb2bacecd19679b6575ae33490dd36
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:518b9be7f0f5648b8f847c7dbb6590d2aa9c7f511d5c40e153ea7270fdbbd441
+size 19677184
diff --git a/permanent/step_100000/l2_ema/checkpoint_100000/_CHECKPOINT_METADATA b/permanent/step_100000/l2_ema/checkpoint_100000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..7d3256ec6351ecbe88f8a40f1d5007cfc9f45337
--- /dev/null
+++ b/permanent/step_100000/l2_ema/checkpoint_100000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778541786361133802, "commit_timestamp_nsecs": 1778541786380448397, "custom_metadata": {}}
\ No newline at end of file
diff --git a/permanent/step_100000/l2_ema/checkpoint_100000/_METADATA b/permanent/step_100000/l2_ema/checkpoint_100000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..ada3cd782e85f12ef9574e31944245c2e7aefd59
--- /dev/null
+++ b/permanent/step_100000/l2_ema/checkpoint_100000/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('l2_ema',)": {"key_metadata": [{"key": "l2_ema", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [29, 5, 256]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/permanent/step_100000/l2_ema/checkpoint_100000/_sharding b/permanent/step_100000/l2_ema/checkpoint_100000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..ac31b218b77c613e621eacb64ab85a4c534b1351
--- /dev/null
+++ b/permanent/step_100000/l2_ema/checkpoint_100000/_sharding
@@ -0,0 +1 @@
+{"bDJfZW1h":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}"}
\ No newline at end of file
diff --git a/permanent/step_100000/l2_ema/checkpoint_100000/d/4d1ab09e48b4d07de7818ddbe5f3c8c4 b/permanent/step_100000/l2_ema/checkpoint_100000/d/4d1ab09e48b4d07de7818ddbe5f3c8c4
new file mode 100644
index 0000000000000000000000000000000000000000..6814fdd146ab082b68af098da9b00d429c6fad81
Binary files /dev/null and b/permanent/step_100000/l2_ema/checkpoint_100000/d/4d1ab09e48b4d07de7818ddbe5f3c8c4 differ
diff --git a/permanent/step_100000/l2_ema/checkpoint_100000/manifest.ocdbt b/permanent/step_100000/l2_ema/checkpoint_100000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..94091269f9a3215d9e03691b869f24e89aeaedc9
Binary files /dev/null and b/permanent/step_100000/l2_ema/checkpoint_100000/manifest.ocdbt differ
diff --git a/permanent/step_100000/l2_ema/checkpoint_100000/ocdbt.process_0/d/c24b85aaaaa34340f6c00cd34b530ca8 b/permanent/step_100000/l2_ema/checkpoint_100000/ocdbt.process_0/d/c24b85aaaaa34340f6c00cd34b530ca8
new file mode 100644
index 0000000000000000000000000000000000000000..6814fdd146ab082b68af098da9b00d429c6fad81
Binary files /dev/null and b/permanent/step_100000/l2_ema/checkpoint_100000/ocdbt.process_0/d/c24b85aaaaa34340f6c00cd34b530ca8 differ
diff --git a/permanent/step_100000/opt_state/checkpoint_100000/ocdbt.process_0/d/67d4789f8c2b945a3abb6a69f2311204 b/permanent/step_100000/opt_state/checkpoint_100000/ocdbt.process_0/d/67d4789f8c2b945a3abb6a69f2311204
new file mode 100644
index 0000000000000000000000000000000000000000..46ac4c085ce436fb082ffef264313140850a745b
--- /dev/null
+++ b/permanent/step_100000/opt_state/checkpoint_100000/ocdbt.process_0/d/67d4789f8c2b945a3abb6a69f2311204
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:721bb65baa35011ec50d618f3b54ead010f2541e0a06ff5bfc65ce1b4d3e0cf4
+size 334404
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/_CHECKPOINT_METADATA b/permanent/step_100000/predictor_ema/checkpoint_100000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..ef4314d548c65bfe810c706b2691d1e387bb782c
--- /dev/null
+++ b/permanent/step_100000/predictor_ema/checkpoint_100000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778541785650324154, "commit_timestamp_nsecs": 1778541786358813067, "custom_metadata": {}}
\ No newline at end of file
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/_METADATA b/permanent/step_100000/predictor_ema/checkpoint_100000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..98e40c42a0223d39dbd11e0b3d709b14663fc8bf
--- /dev/null
+++ b/permanent/step_100000/predictor_ema/checkpoint_100000/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('blocks_0', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_0', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_0', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_0', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_0', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_0', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_0', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_0', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_0', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_0', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_0', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_0', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_0", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_1', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_1', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_1', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_1', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_1', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_1', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_1', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_1', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_1', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_1', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_1', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_1', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_1", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_10', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_10', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_10', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_10', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_10', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_10', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_10', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_10', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_10', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_10', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_10', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_10', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_10", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_11', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_11', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_11', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_11', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('blocks_11', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_11', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_11', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_11', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_11', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_11', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_11', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_11', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_11', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_11', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_11', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_11', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_11', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_11', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_11', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_11', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_11', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_11', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_11", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_2', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_2', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_2', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_2', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_2', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_2', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_2', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_2', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_2', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_2', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_2', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_2', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_2", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_3', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_3', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_3', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_3', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('blocks_3', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_3', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_3', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_3', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_3', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_3', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_3', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_3', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_3', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_3', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_3', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_3', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_3', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_3', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_3', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_3', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_3', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_3', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_3", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_4', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_4', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_4', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_4', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_4', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_4', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_4', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_4', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_4', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_4', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_4', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_4', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_4", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_5', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_5', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_5', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_5', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_5', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_5', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_5', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_5', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_5', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_5', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_5', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_5', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_5", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_6', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_6', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_6', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_6', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_6', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_6', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_6', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_6', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_6', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_6', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_6', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_6', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_6", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_7', 'attn', 'key', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_7', 'attn', 'key', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_7', 'attn', 'out', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_7', 'attn', 'out', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60, 480]}}, "('blocks_7', 'attn', 'query', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_7', 'attn', 'query', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_7', 'attn', 'value', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "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, 60]}}, "('blocks_7', 'attn', 'value', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 8, 60]}}, "('blocks_7', 'attn_adaln', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_7', 'attn_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "attn_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_7', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_7', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_7', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_7', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_7', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_7', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_7', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_7', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_7', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_7', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_7', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_7', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_7", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_8', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_8', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_8', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_8', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_8', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_8', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_8', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_8', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_8', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_8', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_8', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_8', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_8", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('blocks_9', 'conv_adaln', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_9', 'conv_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "conv_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_9', 'dwconv', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_9', 'dwconv', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 480]}}, "('blocks_9', 'mlp_adaln', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1440]}}, "('blocks_9', 'mlp_adaln', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_adaln", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1440]}}, "('blocks_9', 'mlp_fc1', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920]}}, "('blocks_9', 'mlp_fc1', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1920]}}, "('blocks_9', 'mlp_fc2', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_9', 'mlp_fc2', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "mlp_fc2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1920, 480]}}, "('blocks_9', 'pwconv', 'bias')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('blocks_9', 'pwconv', 'kernel')": {"key_metadata": [{"key": "blocks_9", "key_type": 2}, {"key": "pwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 480]}}, "('class_input_embed', 'embedding')": {"key_metadata": [{"key": "class_input_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1001, 1152]}}, "('cond_delta_embed', 'embedding')": {"key_metadata": [{"key": "cond_delta_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [28, 32]}}, "('cond_layer_embed', 'embedding')": {"key_metadata": [{"key": "cond_layer_embed", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [29, 32]}}, "('cond_out', 'bias')": {"key_metadata": [{"key": "cond_out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('cond_out', 'kernel')": {"key_metadata": [{"key": "cond_out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('cond_t_proj', 'bias')": {"key_metadata": [{"key": "cond_t_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('cond_t_proj', 'kernel')": {"key_metadata": [{"key": "cond_t_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 32]}}, "('final_ln', 'bias')": {"key_metadata": [{"key": "final_ln", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('final_ln', 'scale')": {"key_metadata": [{"key": "final_ln", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('gamma_out',)": {"key_metadata": [{"key": "gamma_out", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('in_proj', 'bias')": {"key_metadata": [{"key": "in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480]}}, "('in_proj', 'kernel')": {"key_metadata": [{"key": "in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 480]}}, "('mag_head', 'adaln_mod', 'bias')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [964]}}, "('mag_head', 'adaln_mod', 'kernel')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "adaln_mod", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 964]}}, "('mag_head', 'dwconv', 'bias')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482]}}, "('mag_head', 'dwconv', 'kernel')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "dwconv", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3, 1, 482]}}, "('mag_head', 'pw1', 'bias')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120]}}, "('mag_head', 'pw1', 'kernel')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "pw1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [482, 120]}}, "('mag_head', 'pw2', 'bias')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('mag_head', 'pw2', 'kernel')": {"key_metadata": [{"key": "mag_head", "key_type": 2}, {"key": "pw2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [120, 1]}}, "('out_proj', 'bias')": {"key_metadata": [{"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152]}}, "('out_proj', 'kernel')": {"key_metadata": [{"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [480, 1152]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/_sharding b/permanent/step_100000/predictor_ema/checkpoint_100000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..e953c1dfb239a44e96eeebcd20a1eb878636ca2c
--- /dev/null
+++ b/permanent/step_100000/predictor_ema/checkpoint_100000/_sharding
@@ -0,0 +1 @@
+{"Y29uZF90X3Byb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF90X3Byb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF9kZWx0YV9lbWJlZC5lbWJlZGRpbmc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF9sYXllcl9lbWJlZC5lbWJlZGRpbmc=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF9vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y29uZF9vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Y2xhc3NfaW5wdXRfZW1iZWQuZW1iZWRkaW5n":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzAucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzEwLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9mYzEuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9mYzEua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9mYzIuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLm1scF9mYzIua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ua2V5LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ub3V0LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ucXVlcnkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4ucXVlcnkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4udmFsdWUuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG4udmFsdWUua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG5fYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmF0dG5fYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmNvbnZfYWRhbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmNvbnZfYWRhbG4ua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmR3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLmR3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLnB3Y29udi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzExLnB3Y29udi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzIucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuYXR0bl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzMucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzQucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzUucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzYucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi52YWx1ZS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5rZXkuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5rZXkua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5vdXQuYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5vdXQua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5xdWVyeS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bi5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuYXR0bl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzcucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzguY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzguY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzguZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzguZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzgucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkuY29udl9hZGFsbi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkuY29udl9hZGFsbi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2FkYWxuLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2FkYWxuLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2ZjMS5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2ZjMS5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2ZjMi5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkubWxwX2ZjMi5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkucHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","YmxvY2tzXzkucHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","Z2FtbWFfb3V0":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","ZmluYWxfbG4uYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","ZmluYWxfbG4uc2NhbGU=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5fcHJvai5iaWFz":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","aW5fcHJvai5rZXJuZWw=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","b3V0X3Byb2ouYmlhcw==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","b3V0X3Byb2oua2VybmVs":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQuYWRhbG5fbW9kLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQuYWRhbG5fbW9kLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQuZHdjb252LmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQuZHdjb252Lmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQucHcxLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQucHcxLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQucHcyLmJpYXM=":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}","bWFnX2hlYWQucHcyLmtlcm5lbA==":"{\"sharding_type\": \"SingleDeviceSharding\", \"device_str\": \"TPU_0(process=0,(0,0,0,0))\"}"}
\ No newline at end of file
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/array_metadatas/process_0 b/permanent/step_100000/predictor_ema/checkpoint_100000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..16aca7b61e88e945fe42c867d09c82f96f78d5d0
--- /dev/null
+++ b/permanent/step_100000/predictor_ema/checkpoint_100000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "blocks_0.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_0.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_1.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_10.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_11.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_2.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_3.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_4.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_5.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_6.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.key.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.key.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.out.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.out.kernel", "write_shape": [8, 60, 480], "chunk_shape": [8, 60, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.query.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.query.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.value.bias", "write_shape": [8, 60], "chunk_shape": [8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn.value.kernel", "write_shape": [480, 8, 60], "chunk_shape": [480, 8, 60], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.attn_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_7.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_8.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.conv_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.conv_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.dwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.dwconv.kernel", "write_shape": [3, 3, 1, 480], "chunk_shape": [3, 3, 1, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_adaln.bias", "write_shape": [1440], "chunk_shape": [1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_adaln.kernel", "write_shape": [32, 1440], "chunk_shape": [32, 1440], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_fc1.bias", "write_shape": [1920], "chunk_shape": [1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_fc1.kernel", "write_shape": [480, 1920], "chunk_shape": [480, 1920], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_fc2.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.mlp_fc2.kernel", "write_shape": [1920, 480], "chunk_shape": [1920, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.pwconv.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "blocks_9.pwconv.kernel", "write_shape": [480, 480], "chunk_shape": [480, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "class_input_embed.embedding", "write_shape": [1001, 1152], "chunk_shape": [1001, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_delta_embed.embedding", "write_shape": [28, 32], "chunk_shape": [28, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_layer_embed.embedding", "write_shape": [29, 32], "chunk_shape": [29, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_out.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_out.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_t_proj.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_t_proj.kernel", "write_shape": [1152, 32], "chunk_shape": [1152, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "final_ln.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "final_ln.scale", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "gamma_out", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "in_proj.bias", "write_shape": [480], "chunk_shape": [480], "ext_metadata": null}}, {"array_metadata": {"param_name": "in_proj.kernel", "write_shape": [1152, 480], "chunk_shape": [1152, 480], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.adaln_mod.bias", "write_shape": [964], "chunk_shape": [964], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.adaln_mod.kernel", "write_shape": [32, 964], "chunk_shape": [32, 964], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.dwconv.bias", "write_shape": [482], "chunk_shape": [482], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.dwconv.kernel", "write_shape": [3, 3, 1, 482], "chunk_shape": [3, 3, 1, 482], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.pw1.bias", "write_shape": [120], "chunk_shape": [120], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.pw1.kernel", "write_shape": [482, 120], "chunk_shape": [482, 120], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.pw2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "mag_head.pw2.kernel", "write_shape": [120, 1], "chunk_shape": [120, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "out_proj.bias", "write_shape": [1152], "chunk_shape": [1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "out_proj.kernel", "write_shape": [480, 1152], "chunk_shape": [480, 1152], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/d/02a0891b10ce94012a6320beaa76927a b/permanent/step_100000/predictor_ema/checkpoint_100000/d/02a0891b10ce94012a6320beaa76927a
new file mode 100644
index 0000000000000000000000000000000000000000..1a58257aac4fd5c7e20e3a57c471642a25da9443
Binary files /dev/null and b/permanent/step_100000/predictor_ema/checkpoint_100000/d/02a0891b10ce94012a6320beaa76927a differ
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/24733009192379f8f59f98e1c9ddf6ad b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/24733009192379f8f59f98e1c9ddf6ad
new file mode 100644
index 0000000000000000000000000000000000000000..9be8d529bba13a9207892237858798a7ed5ed719
--- /dev/null
+++ b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/24733009192379f8f59f98e1c9ddf6ad
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cde152cdae7927b70462f11836233f7d87cba71fe6105dbb6efbdc797edb46f
+size 13139968
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/4535724fe7708c015e5ce5eb473d37a7 b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/4535724fe7708c015e5ce5eb473d37a7
new file mode 100644
index 0000000000000000000000000000000000000000..784934bbe677085651552dc96344d9b5227787f9
--- /dev/null
+++ b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/4535724fe7708c015e5ce5eb473d37a7
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:420038afb7054f93d3a772c3125a69a4760c5867546e66a723b610d2dfaa3945
+size 174378
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/8cdf871e48db37b61b685483ab8e907f b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/8cdf871e48db37b61b685483ab8e907f
new file mode 100644
index 0000000000000000000000000000000000000000..7c91c0f5b3f0d61f67a37d041d19db984897a185
Binary files /dev/null and b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/8cdf871e48db37b61b685483ab8e907f differ
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/99be9958200bf1b41a1dbbc186832b4d b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/99be9958200bf1b41a1dbbc186832b4d
new file mode 100644
index 0000000000000000000000000000000000000000..a4ea8d5d86efdd2a0fb099951bd3bcc5f3f4f1d8
--- /dev/null
+++ b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/99be9958200bf1b41a1dbbc186832b4d
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff6ede3d7c81c8c8b89ca7a0d76cc335c3cb4b2916edd1dc135e16423f629eb0
+size 103137280
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/9ce36b1e698628593f883e93ca0b483b b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/9ce36b1e698628593f883e93ca0b483b
new file mode 100644
index 0000000000000000000000000000000000000000..a56fb41671f713da7e2b320c6bd3e411e7b569fe
Binary files /dev/null and b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/9ce36b1e698628593f883e93ca0b483b differ
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/c2bb4d6d8c856a81e530270f90dc3015 b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/c2bb4d6d8c856a81e530270f90dc3015
new file mode 100644
index 0000000000000000000000000000000000000000..b8a71b1a55313afc5bd6c24e885e03d9c416ca0a
Binary files /dev/null and b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/c2bb4d6d8c856a81e530270f90dc3015 differ
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/db2166161a485e2376576a8435feefa6 b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/db2166161a485e2376576a8435feefa6
new file mode 100644
index 0000000000000000000000000000000000000000..849a6805ff0c202cd6aa4e7320ba28dab061ef8b
Binary files /dev/null and b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/d/db2166161a485e2376576a8435feefa6 differ
diff --git a/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/manifest.ocdbt b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..4c9510e6b1a7f3c35320b3df7e67a60f9ec56066
Binary files /dev/null and b/permanent/step_100000/predictor_ema/checkpoint_100000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/permanent/step_200000/checkpoint_200000/ocdbt.process_0/d/73b277c20e1a8d9b9cacf922bfe6c166 b/permanent/step_200000/checkpoint_200000/ocdbt.process_0/d/73b277c20e1a8d9b9cacf922bfe6c166
new file mode 100644
index 0000000000000000000000000000000000000000..f45a53982c4261cd4e33b7a0eb1b346db1b8672e
Binary files /dev/null and b/permanent/step_200000/checkpoint_200000/ocdbt.process_0/d/73b277c20e1a8d9b9cacf922bfe6c166 differ
diff --git a/permanent/step_200000/checkpoint_200000/ocdbt.process_0/d/b2740f6c6449f1c164c9bc7918fe176e b/permanent/step_200000/checkpoint_200000/ocdbt.process_0/d/b2740f6c6449f1c164c9bc7918fe176e
new file mode 100644
index 0000000000000000000000000000000000000000..c86e1153e9284a354cdc996e1d5b9cc660cd8c02
--- /dev/null
+++ b/permanent/step_200000/checkpoint_200000/ocdbt.process_0/d/b2740f6c6449f1c164c9bc7918fe176e
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2285e999128fbac13611c826ab80b03411e961424be5c54d6d453bcfcc97f89
+size 181764
diff --git a/permanent/step_200000/ema/checkpoint_200000/d/b49eab8fd7c555311dadcd52467efd06 b/permanent/step_200000/ema/checkpoint_200000/d/b49eab8fd7c555311dadcd52467efd06
new file mode 100644
index 0000000000000000000000000000000000000000..086d768c4391c2ccbe58f1b006e74254dfcee115
Binary files /dev/null and b/permanent/step_200000/ema/checkpoint_200000/d/b49eab8fd7c555311dadcd52467efd06 differ
diff --git a/permanent/step_200000/ema/checkpoint_200000/ocdbt.process_0/d/95af092583d380070ceacbee12c385fd b/permanent/step_200000/ema/checkpoint_200000/ocdbt.process_0/d/95af092583d380070ceacbee12c385fd
new file mode 100644
index 0000000000000000000000000000000000000000..b76a321f40500fa32bd90b8166d189a1db365cbf
--- /dev/null
+++ b/permanent/step_200000/ema/checkpoint_200000/ocdbt.process_0/d/95af092583d380070ceacbee12c385fd
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f274ecd64783076742d15232573ff0d0ca0de8b99764412bc03daa2614a185f0
+size 364669
diff --git a/permanent/step_200000/opt_state/checkpoint_200000/d/c30307ac69dd6022357c623120856fe5 b/permanent/step_200000/opt_state/checkpoint_200000/d/c30307ac69dd6022357c623120856fe5
new file mode 100644
index 0000000000000000000000000000000000000000..ac88f19e7e117f9d2690245496af00c48301c1d4
Binary files /dev/null and b/permanent/step_200000/opt_state/checkpoint_200000/d/c30307ac69dd6022357c623120856fe5 differ
diff --git a/permanent/step_200000/opt_state/checkpoint_200000/ocdbt.process_0/d/c827e39babc608b717bc851428d44939 b/permanent/step_200000/opt_state/checkpoint_200000/ocdbt.process_0/d/c827e39babc608b717bc851428d44939
new file mode 100644
index 0000000000000000000000000000000000000000..3f571300bfc2e42232b196d5f5f5f901d478d9a4
--- /dev/null
+++ b/permanent/step_200000/opt_state/checkpoint_200000/ocdbt.process_0/d/c827e39babc608b717bc851428d44939
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a680ca3d236b32331a360bd2ab53330bd4b1314b645653cd8eab49e45223478
+size 286788
diff --git a/permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/380c55ccd750470f7dadc2175c156a5b b/permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/380c55ccd750470f7dadc2175c156a5b
new file mode 100644
index 0000000000000000000000000000000000000000..9fbdeda714595342b2a97a4770334c6254760d9a
--- /dev/null
+++ b/permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/380c55ccd750470f7dadc2175c156a5b
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:961fb5146dae68c90eff3853c7a837e38eab2ab8b3d2909eafcdab4fd5950858
+size 105504768
diff --git a/permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/3e028d2bbb4b009cc0a1652900d6dbae b/permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/3e028d2bbb4b009cc0a1652900d6dbae
new file mode 100644
index 0000000000000000000000000000000000000000..2c7ee6edf9cc3ed513349a453643bf987a4d9426
--- /dev/null
+++ b/permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/3e028d2bbb4b009cc0a1652900d6dbae
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c888fe21126ef65e776d111f8f638cee9f27b64192d99e1d2bf07c64db407e5
+size 10805248
diff --git a/permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/d7fd7d3fdd1c351adac17ed11a0323e3 b/permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/d7fd7d3fdd1c351adac17ed11a0323e3
new file mode 100644
index 0000000000000000000000000000000000000000..883605c49342a083391cd55de883fc178c099751
--- /dev/null
+++ b/permanent/step_200000/predictor_ema/checkpoint_200000/ocdbt.process_0/d/d7fd7d3fdd1c351adac17ed11a0323e3
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8552391817d7ab9ea1bc1a7b66baf0697be0ef5c48c2c045ce9ebd1f0310e11
+size 206375
diff --git a/permanent/step_400000/checkpoint_400000/ocdbt.process_0/d/058cbf30ed1dce8e04ea12f92cc3522e b/permanent/step_400000/checkpoint_400000/ocdbt.process_0/d/058cbf30ed1dce8e04ea12f92cc3522e
new file mode 100644
index 0000000000000000000000000000000000000000..3a3f08723a5e633cb7a531bb4c68be2d21cfdced
--- /dev/null
+++ b/permanent/step_400000/checkpoint_400000/ocdbt.process_0/d/058cbf30ed1dce8e04ea12f92cc3522e
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0319d3c139fa7b886e41a2e61016b80053a15b6d2095a88a110a81d31d0702c8
+size 553379
diff --git a/permanent/step_400000/ema/checkpoint_400000/ocdbt.process_0/d/910e1420ad62226535ee916a97e440f5 b/permanent/step_400000/ema/checkpoint_400000/ocdbt.process_0/d/910e1420ad62226535ee916a97e440f5
new file mode 100644
index 0000000000000000000000000000000000000000..d50625080f4e77fe66c8788b6d88bdec64f3c93f
--- /dev/null
+++ b/permanent/step_400000/ema/checkpoint_400000/ocdbt.process_0/d/910e1420ad62226535ee916a97e440f5
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b545fc76f97bbc35ec4d761a778d622651ddb918404eeaab39e33337f5b7d1b
+size 4956160
diff --git a/permanent/step_400000/opt_state/checkpoint_400000/ocdbt.process_0/d/b6c8e12debdd19d0beeece616d43dcf5 b/permanent/step_400000/opt_state/checkpoint_400000/ocdbt.process_0/d/b6c8e12debdd19d0beeece616d43dcf5
new file mode 100644
index 0000000000000000000000000000000000000000..b9244c448c347762c93e49f63ba0c5dbafffccb7
--- /dev/null
+++ b/permanent/step_400000/opt_state/checkpoint_400000/ocdbt.process_0/d/b6c8e12debdd19d0beeece616d43dcf5
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:533eb87f187e21aea99d8134a1be55c9e92b8771044b0e432ad86d48c7987b25
+size 451062
diff --git a/permanent/step_400000/opt_state/checkpoint_400000/ocdbt.process_0/d/d53aa29264c20ee8d87912444fba44d1 b/permanent/step_400000/opt_state/checkpoint_400000/ocdbt.process_0/d/d53aa29264c20ee8d87912444fba44d1
new file mode 100644
index 0000000000000000000000000000000000000000..f9926629f6fb2cbdf694fffb4e55b57ecd7e3be4
--- /dev/null
+++ b/permanent/step_400000/opt_state/checkpoint_400000/ocdbt.process_0/d/d53aa29264c20ee8d87912444fba44d1
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb99f3c6a781998d27cb4d30e25ffea8447fbfd7defe7dcb37b19083aebcd551
+size 169936
diff --git a/permanent/step_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/4c4188133a37300d5ada5ac8a3105934 b/permanent/step_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/4c4188133a37300d5ada5ac8a3105934
new file mode 100644
index 0000000000000000000000000000000000000000..bc1a4180cad628855bd9afc129d0c69e231810a0
--- /dev/null
+++ b/permanent/step_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/4c4188133a37300d5ada5ac8a3105934
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c0aff8e0d29ec62f7861204923ba432378e7efe286a7bdcf061bff14a1c606b
+size 112230400
diff --git a/permanent/step_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/b399e969b4326f5b52f8e76b7cc2ae7c b/permanent/step_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/b399e969b4326f5b52f8e76b7cc2ae7c
new file mode 100644
index 0000000000000000000000000000000000000000..fca6b45dc7df3b37ff7affddd1a3d1de3a5a2908
--- /dev/null
+++ b/permanent/step_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/b399e969b4326f5b52f8e76b7cc2ae7c
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1344e7693e3484a87a44b917d0ffa37cc485c47453ef7354172980756b93d46e
+size 4268032
diff --git a/wandb_run_summary.json b/wandb_run_summary.json
new file mode 100644
index 0000000000000000000000000000000000000000..420fe21790d4a17717993c28fdde1a0f6554a390
--- /dev/null
+++ b/wandb_run_summary.json
@@ -0,0 +1,502 @@
+{
+  "best": {
+    "train/loss": {
+      "step": 361000,
+      "value": 0.7244978547096252
+    },
+    "train/loss_gen": {
+      "step": 361000,
+      "value": 0.7135342955589294
+    },
+    "train/loss_total": {
+      "step": 361000,
+      "value": 0.7244978547096252
+    },
+    "val/FID": {
+      "step": 400000,
+      "value": 10.443120749429092
+    },
+    "val/InceptionScore": {
+      "step": 400000,
+      "value": 86.12591147913444
+    },
+    "val/Precision": {
+      "step": 200000,
+      "value": 0.4954
+    },
+    "val/Recall": {
+      "step": 400000,
+      "value": 0.574
+    },
+    "val/sFID": {
+      "step": 400000,
+      "value": 17.367982137214973
+    }
+  },
+  "config": {
+    "batch_size": 256,
+    "block_corr_batches": 2,
+    "block_corr_freq": 0,
+    "cfg_dropout_rate": 0.1,
+    "ckpt_dir": "/home/nguyenthanhlam/checkpoints/ablate-xl-magcalib",
+    "ckpt_keep_steps": "100000,200000,400000",
+    "ckpt_latest_freq": 20000,
+    "ckpt_verify_step": 0,
+    "data_path": "/home/nguyenthanhlam/kaggle_downloads/datasets/train",
+    "direct_allow_identity_pairs": false,
+    "direct_joint_pairs": 1,
+    "direct_num_pairs": 1,
+    "direct_pair_mode": "trunc_normal_centered",
+    "direct_pair_order_mode": "forward",
+    "direct_predictor_only_pairs": 0,
+    "ema_decay": 0.9999,
+    "epochs": 400,
+    "eval_batches": 1,
+    "eval_freq": 20000,
+    "fid_batch_size": 256,
+    "fid_cfg_scale": 1.0,
+    "fid_eval_local_batch": 32,
+    "fid_freq": 10000,
+    "fid_num_steps": 250,
+    "fid_skip_eval": false,
+    "fid_skip_timestep_mode": "all",
+    "fid_steps": "50000,100000,200000,400000",
+    "grad_clip": 1.0,
+    "inception_score": true,
+    "inception_score_splits": 10,
+    "inception_score_weights": "/home/nguyenthanhlam/models/inception_v3_google-0cc3c7bd.pth",
+    "lambda_output_distill": 0.05,
+    "lambda_private": 1.0,
+    "learning_rate": 0.0001,
+    "linear_probe": false,
+    "log_freq": 1000,
+    "mock_data": false,
+    "model_size": "XL",
+    "no_official_training_overrides": false,
+    "no_wandb": false,
+    "num_fid_samples": 50000,
+    "official_training_mode": false,
+    "output_distill": true,
+    "output_distill_every": 1,
+    "output_distill_full_backbone_start_step": 0,
+    "output_distill_pair_mode": "trunc_normal_centered",
+    "output_distill_ratio": 0.1,
+    "output_distill_target_mode": "model_output",
+    "output_distill_update_mode": "predictor_plus_all",
+    "pair_center_loc": null,
+    "pair_center_sigma": 4.7,
+    "pair_uniform_anneal_start_step": 0,
+    "pair_uniform_anneal_steps": 100000,
+    "pr_full_mode": false,
+    "pr_k": 3,
+    "pr_max_samples": 5000,
+    "precision_recall": true,
+    "preflight_checks": false,
+    "preflight_fid_memory_probe": false,
+    "preflight_fid_samples": 16,
+    "preflight_only": false,
+    "preflight_sample_count": 4,
+    "private_cosine_mode": "bnd",
+    "private_loss": true,
+    "private_max_pairs": 4,
+    "private_pair_mode": "random",
+    "private_start_step": 0,
+    "private_use_residual": true,
+    "private_warmup_iters": 0,
+    "probe_eval_batches": 4,
+    "probe_layer": null,
+    "probe_save_path": null,
+    "resume": true,
+    "sample_cfg_scale": 1.0,
+    "sample_freq": 0,
+    "sample_num_steps": 50,
+    "sample_skip_eval": false,
+    "shortcut_activation_huber_delta": 1.0,
+    "shortcut_bootstrap_detach_source": true,
+    "shortcut_bootstrap_reverse_stopgrad": false,
+    "shortcut_debug_gap_log_freq": 10000,
+    "shortcut_debug_gap_logs": false,
+    "shortcut_l2_ema_alpha": 0.01,
+    "shortcut_lambda_boot": 0.25,
+    "shortcut_lambda_boot_mag": 0.1875,
+    "shortcut_lambda_dir": 1.0,
+    "shortcut_lambda_mag": 0.375,
+    "shortcut_lambda_skip_fm": 0.0,
+    "shortcut_loss_mode": "direction_magnitude",
+    "shortcut_mag_abs_center": 4.0,
+    "shortcut_mag_abs_scale": 1.0,
+    "shortcut_mag_ce_bins": 33,
+    "shortcut_mag_ce_temperature": 1.0,
+    "shortcut_mag_clip_max": 6.0,
+    "shortcut_mag_clip_min": 3.0,
+    "shortcut_mag_grad_mode": "full",
+    "shortcut_mag_loss_mode": "mse",
+    "shortcut_mag_reg_weight": 1.0,
+    "shortcut_mag_scale": 4.0,
+    "shortcut_predictor": "hybrid_depth30m",
+    "shortcut_predictor_adaln_zero": null,
+    "shortcut_predictor_arch": "existing",
+    "shortcut_predictor_attention_every": null,
+    "shortcut_predictor_class_fusion": "add",
+    "shortcut_predictor_depth": 12,
+    "shortcut_predictor_dilation_cycle": "1,2,4",
+    "shortcut_predictor_ema_decay": 0.999,
+    "shortcut_predictor_grad_clip": 1.0,
+    "shortcut_predictor_grid_size": null,
+    "shortcut_predictor_hidden_size": 480,
+    "shortcut_predictor_lr": 0.0001,
+    "shortcut_predictor_mlp_ratio": null,
+    "shortcut_predictor_normalize_input": false,
+    "shortcut_predictor_num_heads": 8,
+    "shortcut_predictor_residual_output": null,
+    "shortcut_predictor_use_class_input": true,
+    "shortcut_predictor_use_timestep": true,
+    "shortcut_predictor_weight_decay": 0.1,
+    "shortcut_skip_in_loop_detach_source": true,
+    "shortcut_skip_in_loop_gap": 2,
+    "shortcut_skip_in_loop_gap_loc": 7.0,
+    "shortcut_skip_in_loop_gap_mode": "truncated-normal",
+    "shortcut_skip_in_loop_gap_sigma": 4.7,
+    "shortcut_skip_in_loop_max_gap": 23,
+    "shortcut_skip_in_loop_prob": 0.0,
+    "shortcut_skip_in_loop_warmup_steps": 5000,
+    "shortcut_timesteps": 50,
+    "shortcut_training_mode": "direction-magnitude",
+    "steps_per_epoch": 1000,
+    "timestep_logit_mean": 0.0,
+    "timestep_logit_std": 1.0,
+    "timestep_sampling_mode": "logit_normal",
+    "uniform_transition_diag_freq": 1000,
+    "uniform_transition_diag_logs": false,
+    "vae_decode_batch_size": 256,
+    "vae_hf_config": "/kaggle/working/huggingface_cache/hub/models--stabilityai--sd-vae-ft-ema/snapshots/f04b2c4b98319346dad8c65879f680b1997b204a/config.json",
+    "vae_model": "/home/nguyenthanhlam/models/sdvae-ema-flax-default-1",
+    "val_data_path": "/home/nguyenthanhlam/kaggle_downloads/datasets/val",
+    "wandb_project": "selfflow-jax",
+    "weight_decay": 0.1
+  },
+  "latest": {
+    "train": {
+      "perf/accumulated_train_tflops": 31560794.080500036,
+      "perf/train_step_tflops": 255.32552333580168,
+      "perf/train_step_time": 0.34336117792129517,
+      "train/boot_activation_pred_rms": 0.0,
+      "train/boot_activation_rel_residual_rms": 0.0,
+      "train/boot_activation_target_rms": 0.0,
+      "train/bootstrap_reverse_stopgrad": 0.0,
+      "train/common_norm": 13.553272247314453,
+      "train/cos_boot": 0.9971438646316528,
+      "train/cos_dir": 0.9848425984382629,
+      "train/debug_gap_1/cos_dir": 0.0,
+      "train/debug_gap_1/delta_m_mae": 0.0,
+      "train/debug_gap_1/loss_dir": 0.0,
+      "train/debug_gap_1/loss_mag": 0.0,
+      "train/debug_gap_10/cos_dir": 0.0,
+      "train/debug_gap_10/delta_m_mae": 0.0,
+      "train/debug_gap_10/loss_dir": 0.0,
+      "train/debug_gap_10/loss_mag": 0.0,
+      "train/debug_gap_2/cos_dir": 0.0,
+      "train/debug_gap_2/delta_m_mae": 0.0,
+      "train/debug_gap_2/loss_dir": 0.0,
+      "train/debug_gap_2/loss_mag": 0.0,
+      "train/debug_gap_3/cos_dir": 0.0,
+      "train/debug_gap_3/delta_m_mae": 0.0,
+      "train/debug_gap_3/loss_dir": 0.0,
+      "train/debug_gap_3/loss_mag": 0.0,
+      "train/debug_gap_4/cos_dir": 0.0,
+      "train/debug_gap_4/delta_m_mae": 0.0,
+      "train/debug_gap_4/loss_dir": 0.0,
+      "train/debug_gap_4/loss_mag": 0.0,
+      "train/debug_gap_5/cos_dir": 0.0,
+      "train/debug_gap_5/delta_m_mae": 0.0,
+      "train/debug_gap_5/loss_dir": 0.0,
+      "train/debug_gap_5/loss_mag": 0.0,
+      "train/debug_gap_6/cos_dir": 0.0,
+      "train/debug_gap_6/delta_m_mae": 0.0,
+      "train/debug_gap_6/loss_dir": 0.0,
+      "train/debug_gap_6/loss_mag": 0.0,
+      "train/debug_gap_7/cos_dir": 0.0,
+      "train/debug_gap_7/delta_m_mae": 0.0,
+      "train/debug_gap_7/loss_dir": 0.0,
+      "train/debug_gap_7/loss_mag": 0.0,
+      "train/debug_gap_8/cos_dir": 0.0,
+      "train/debug_gap_8/delta_m_mae": 0.0,
+      "train/debug_gap_8/loss_dir": 0.0,
+      "train/debug_gap_8/loss_mag": 0.0,
+      "train/debug_gap_9/cos_dir": 0.0,
+      "train/debug_gap_9/delta_m_mae": 0.0,
+      "train/debug_gap_9/loss_dir": 0.0,
+      "train/debug_gap_9/loss_mag": 0.0,
+      "train/debug_gap_log_freq": 10000.0,
+      "train/debug_gap_logs_active": 0.0,
+      "train/debug_pair_10_12/cos_dir": 0.0,
+      "train/debug_pair_10_12/delta_m_mae": 0.0,
+      "train/debug_pair_10_12/loss_dir": 0.0,
+      "train/debug_pair_10_12/loss_mag": 0.0,
+      "train/debug_pair_1_10/cos_dir": 0.0,
+      "train/debug_pair_1_10/delta_m_mae": 0.0,
+      "train/debug_pair_1_10/loss_dir": 0.0,
+      "train/debug_pair_1_10/loss_mag": 0.0,
+      "train/debug_pair_1_3/cos_dir": 0.0,
+      "train/debug_pair_1_3/delta_m_mae": 0.0,
+      "train/debug_pair_1_3/loss_dir": 0.0,
+      "train/debug_pair_1_3/loss_mag": 0.0,
+      "train/debug_pair_1_5/cos_dir": 0.0,
+      "train/debug_pair_1_5/delta_m_mae": 0.0,
+      "train/debug_pair_1_5/loss_dir": 0.0,
+      "train/debug_pair_1_5/loss_mag": 0.0,
+      "train/debug_pair_3_12/cos_dir": 0.0,
+      "train/debug_pair_3_12/delta_m_mae": 0.0,
+      "train/debug_pair_3_12/loss_dir": 0.0,
+      "train/debug_pair_3_12/loss_mag": 0.0,
+      "train/debug_pair_3_5/cos_dir": 0.0,
+      "train/debug_pair_3_5/delta_m_mae": 0.0,
+      "train/debug_pair_3_5/loss_dir": 0.0,
+      "train/debug_pair_3_5/loss_mag": 0.0,
+      "train/debug_pair_3_7/cos_dir": 0.0,
+      "train/debug_pair_3_7/delta_m_mae": 0.0,
+      "train/debug_pair_3_7/loss_dir": 0.0,
+      "train/debug_pair_3_7/loss_mag": 0.0,
+      "train/debug_pair_5_7/cos_dir": 0.0,
+      "train/debug_pair_5_7/delta_m_mae": 0.0,
+      "train/debug_pair_5_7/loss_dir": 0.0,
+      "train/debug_pair_5_7/loss_mag": 0.0,
+      "train/debug_pair_6_12/cos_dir": 0.0,
+      "train/debug_pair_6_12/delta_m_mae": 0.0,
+      "train/debug_pair_6_12/loss_dir": 0.0,
+      "train/debug_pair_6_12/loss_mag": 0.0,
+      "train/delta_m_clip_rate": 0.0,
+      "train/delta_m_mae": 0.017344729974865913,
+      "train/delta_m_ratio_error": 1.0175893306732178,
+      "train/delta_m_rmse": 0.021340016275644302,
+      "train/direct_activation_pred_rms_joint": 0.0,
+      "train/direct_activation_rel_residual_rms_joint": 0.0,
+      "train/direct_activation_target_rms_joint": 0.0,
+      "train/direct_gap": 10.0,
+      "train/direct_gap_joint": 10.0,
+      "train/direct_gap_ponly_1": 0.0,
+      "train/direct_gap_ponly_2": 0.0,
+      "train/direct_loss_mode": 0.0,
+      "train/direct_num_pairs": 1.0,
+      "train/direct_pair_joint_a": 8.5,
+      "train/direct_pair_joint_b": 18.5,
+      "train/direct_pair_ponly_1_a": -1.0,
+      "train/direct_pair_ponly_1_b": -1.0,
+      "train/direct_pair_ponly_2_a": -1.0,
+      "train/direct_pair_ponly_2_b": -1.0,
+      "train/ema_decay": 0.9998999834060669,
+      "train/grad_norm": 0.09155187010765076,
+      "train/l_private": 0.0011048904852941632,
+      "train/lambda_private_effective": 1.0,
+      "train/loss": 0.7548577785491943,
+      "train/loss_boot": 0.0028561055660247803,
+      "train/loss_boot_aux": 1.7881393432617188e-05,
+      "train/loss_boot_mag": 1.7881393432617188e-05,
+      "train/loss_dir": 0.015157431364059448,
+      "train/loss_dir_joint": 0.015157431364059448,
+      "train/loss_dir_ponly_mean": 0.0,
+      "train/loss_direct_3pair": 0.015163120813667774,
+      "train/loss_direct_aux_joint": 1.5173697647696827e-05,
+      "train/loss_direct_aux_ponly_mean": 0.0,
+      "train/loss_direct_joint": 0.015163120813667774,
+      "train/loss_direct_npairs": 0.015163120813667774,
+      "train/loss_direct_ponly_1": 0.0,
+      "train/loss_direct_ponly_2": 0.0,
+      "train/loss_gen": 0.7374104857444763,
+      "train/loss_mag": 1.5173697647696827e-05,
+      "train/loss_mag_joint": 1.5173697647696827e-05,
+      "train/loss_mag_ponly_mean": 0.0,
+      "train/loss_output_distill": 0.009238230995833874,
+      "train/loss_output_distill_weighted": 0.0004619115497916937,
+      "train/loss_skip_fm": 0.0,
+      "train/loss_total": 0.7548577785491943,
+      "train/lr_backbone": 9.999999747378752e-05,
+      "train/lr_predictor": 9.999999747378752e-05,
+      "train/output_distill_a": 11.25,
+      "train/output_distill_b": 19.5,
+      "train/output_distill_batch_size": 6.0,
+      "train/output_distill_full_backbone_active": 1.0,
+      "train/output_distill_full_backbone_start_step": 0.0,
+      "train/output_distill_gap": 8.25,
+      "train/output_distill_ratio": 0.10000000149011612,
+      "train/output_distill_update_mode": 2.0,
+      "train/pair_center_loc": -1.0,
+      "train/pair_center_sigma": 4.699999809265137,
+      "train/pair_uniform_anneal_start_step": 0.0,
+      "train/pair_uniform_anneal_steps": 100000.0,
+      "train/pair_uniform_mix": 1.0,
+      "train/param_norm": 597.7843627929688,
+      "train/predictor_ema_decay": 0.9990000128746033,
+      "train/predictor_normalize_input": 0.0,
+      "train/predictor_use_class_input": 1.0,
+      "train/predictor_use_timestep": 1.0,
+      "train/private_avg_norm": 8.247620582580566,
+      "train/private_cosine_mode": 0.0,
+      "train/private_pair_mode": 1.0,
+      "train/private_pairwise_cosine": -0.029284359887242317,
+      "train/private_use_residual": 1.0,
+      "train/private_warmup_scale": 1.0,
+      "train/sampled_timestep_q": 24.97265625,
+      "train/sampled_timestep_tau": 0.5096460580825806,
+      "train/shortcut_activation_huber_delta": 1.0,
+      "train/shortcut_loss_mode": 0.0,
+      "train/shortcut_mag_ce_bins": 33.0,
+      "train/shortcut_mag_ce_temperature": 1.0,
+      "train/shortcut_mag_grad_mode": 0.0,
+      "train/shortcut_mag_loss_mode": 0.0,
+      "train/shortcut_mag_loss_stop_backbone": 0.0,
+      "train/shortcut_mag_reg_weight": 1.0,
+      "train/skip_in_loop_do": 0.0,
+      "train/skip_in_loop_gap": 0.0,
+      "train/skip_in_loop_prob": 0.0,
+      "train/skip_in_loop_source": -1.0,
+      "train/skip_in_loop_target": -1.0,
+      "train/step": 400000.0,
+      "train/timestep_logit_mean": 0.0,
+      "train/timestep_logit_std": 1.0,
+      "train/timestep_sampling_mode": 1.0,
+      "train/uniform_transition_diag/all/count": 0.0,
+      "train/uniform_transition_diag/all/l_pred": 0.0,
+      "train/uniform_transition_diag/all/share": 0.0,
+      "train/uniform_transition_diag/freq": 1000.0,
+      "train/uniform_transition_diag/gap_1/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_1/share": 0.0,
+      "train/uniform_transition_diag/gap_10/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_10/share": 0.0,
+      "train/uniform_transition_diag/gap_11/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_11/share": 0.0,
+      "train/uniform_transition_diag/gap_12/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_12/share": 0.0,
+      "train/uniform_transition_diag/gap_13/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_13/share": 0.0,
+      "train/uniform_transition_diag/gap_14/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_14/share": 0.0,
+      "train/uniform_transition_diag/gap_15/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_15/share": 0.0,
+      "train/uniform_transition_diag/gap_16/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_16/share": 0.0,
+      "train/uniform_transition_diag/gap_17/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_17/share": 0.0,
+      "train/uniform_transition_diag/gap_18/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_18/share": 0.0,
+      "train/uniform_transition_diag/gap_19/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_19/share": 0.0,
+      "train/uniform_transition_diag/gap_2/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_2/share": 0.0,
+      "train/uniform_transition_diag/gap_20/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_20/share": 0.0,
+      "train/uniform_transition_diag/gap_21/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_21/share": 0.0,
+      "train/uniform_transition_diag/gap_22/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_22/share": 0.0,
+      "train/uniform_transition_diag/gap_23/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_23/share": 0.0,
+      "train/uniform_transition_diag/gap_24/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_24/share": 0.0,
+      "train/uniform_transition_diag/gap_25/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_25/share": 0.0,
+      "train/uniform_transition_diag/gap_26/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_26/share": 0.0,
+      "train/uniform_transition_diag/gap_27/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_27/share": 0.0,
+      "train/uniform_transition_diag/gap_28/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_28/share": 0.0,
+      "train/uniform_transition_diag/gap_3/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_3/share": 0.0,
+      "train/uniform_transition_diag/gap_4/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_4/share": 0.0,
+      "train/uniform_transition_diag/gap_5/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_5/share": 0.0,
+      "train/uniform_transition_diag/gap_6/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_6/share": 0.0,
+      "train/uniform_transition_diag/gap_7/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_7/share": 0.0,
+      "train/uniform_transition_diag/gap_8/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_8/share": 0.0,
+      "train/uniform_transition_diag/gap_9/l_pred": 0.0,
+      "train/uniform_transition_diag/gap_9/share": 0.0,
+      "train/uniform_transition_diag/source_1/l_pred": 0.0,
+      "train/uniform_transition_diag/source_1/share": 0.0,
+      "train/uniform_transition_diag/source_10/l_pred": 0.0,
+      "train/uniform_transition_diag/source_10/share": 0.0,
+      "train/uniform_transition_diag/source_11/l_pred": 0.0,
+      "train/uniform_transition_diag/source_11/share": 0.0,
+      "train/uniform_transition_diag/source_12/l_pred": 0.0,
+      "train/uniform_transition_diag/source_12/share": 0.0,
+      "train/uniform_transition_diag/source_13/l_pred": 0.0,
+      "train/uniform_transition_diag/source_13/share": 0.0,
+      "train/uniform_transition_diag/source_14/l_pred": 0.0,
+      "train/uniform_transition_diag/source_14/share": 0.0,
+      "train/uniform_transition_diag/source_15/l_pred": 0.0,
+      "train/uniform_transition_diag/source_15/share": 0.0,
+      "train/uniform_transition_diag/source_16/l_pred": 0.0,
+      "train/uniform_transition_diag/source_16/share": 0.0,
+      "train/uniform_transition_diag/source_17/l_pred": 0.0,
+      "train/uniform_transition_diag/source_17/share": 0.0,
+      "train/uniform_transition_diag/source_18/l_pred": 0.0,
+      "train/uniform_transition_diag/source_18/share": 0.0,
+      "train/uniform_transition_diag/source_19/l_pred": 0.0,
+      "train/uniform_transition_diag/source_19/share": 0.0,
+      "train/uniform_transition_diag/source_2/l_pred": 0.0,
+      "train/uniform_transition_diag/source_2/share": 0.0,
+      "train/uniform_transition_diag/source_20/l_pred": 0.0,
+      "train/uniform_transition_diag/source_20/share": 0.0,
+      "train/uniform_transition_diag/source_21/l_pred": 0.0,
+      "train/uniform_transition_diag/source_21/share": 0.0,
+      "train/uniform_transition_diag/source_22/l_pred": 0.0,
+      "train/uniform_transition_diag/source_22/share": 0.0,
+      "train/uniform_transition_diag/source_23/l_pred": 0.0,
+      "train/uniform_transition_diag/source_23/share": 0.0,
+      "train/uniform_transition_diag/source_24/l_pred": 0.0,
+      "train/uniform_transition_diag/source_24/share": 0.0,
+      "train/uniform_transition_diag/source_25/l_pred": 0.0,
+      "train/uniform_transition_diag/source_25/share": 0.0,
+      "train/uniform_transition_diag/source_26/l_pred": 0.0,
+      "train/uniform_transition_diag/source_26/share": 0.0,
+      "train/uniform_transition_diag/source_27/l_pred": 0.0,
+      "train/uniform_transition_diag/source_27/share": 0.0,
+      "train/uniform_transition_diag/source_3/l_pred": 0.0,
+      "train/uniform_transition_diag/source_3/share": 0.0,
+      "train/uniform_transition_diag/source_4/l_pred": 0.0,
+      "train/uniform_transition_diag/source_4/share": 0.0,
+      "train/uniform_transition_diag/source_5/l_pred": 0.0,
+      "train/uniform_transition_diag/source_5/share": 0.0,
+      "train/uniform_transition_diag/source_6/l_pred": 0.0,
+      "train/uniform_transition_diag/source_6/share": 0.0,
+      "train/uniform_transition_diag/source_7/l_pred": 0.0,
+      "train/uniform_transition_diag/source_7/share": 0.0,
+      "train/uniform_transition_diag/source_8/l_pred": 0.0,
+      "train/uniform_transition_diag/source_8/share": 0.0,
+      "train/uniform_transition_diag/source_9/l_pred": 0.0,
+      "train/uniform_transition_diag/source_9/share": 0.0,
+      "train/uniform_transition_diag/timestep_bin_0_high_noise/l_pred": 0.0,
+      "train/uniform_transition_diag/timestep_bin_0_high_noise/share": 0.0,
+      "train/uniform_transition_diag/timestep_bin_1_mid_high_noise/l_pred": 0.0,
+      "train/uniform_transition_diag/timestep_bin_1_mid_high_noise/share": 0.0,
+      "train/uniform_transition_diag/timestep_bin_2_mid_noise/l_pred": 0.0,
+      "train/uniform_transition_diag/timestep_bin_2_mid_noise/share": 0.0,
+      "train/uniform_transition_diag/timestep_bin_3_mid_low_noise/l_pred": 0.0,
+      "train/uniform_transition_diag/timestep_bin_3_mid_low_noise/share": 0.0,
+      "train/uniform_transition_diag/timestep_bin_4_low_noise/l_pred": 0.0,
+      "train/uniform_transition_diag/timestep_bin_4_low_noise/share": 0.0,
+      "train/uniform_transition_diag_active": 0.0,
+      "train/v_abs_mean": 1.0422022342681885,
+      "train/v_pred_abs_mean": 0.78125,
+      "train/y_ab_norm_max": 50.42624282836914,
+      "train/y_ab_norm_mean": 31.953269958496094,
+      "train/y_ab_norm_min": 23.343786239624023
+    },
+    "val": {
+      "train/step": 400000.0,
+      "val/FID": 10.443120749429092,
+      "val/InceptionScore": 86.12591147913444,
+      "val/InceptionScore_std": 2.140928267267062,
+      "val/PR_eval_samples": 5000.0,
+      "val/Precision": 0.4948,
+      "val/Recall": 0.574,
+      "val/sFID": 17.367982137214973
+    }
+  },
+  "latest_step": 400000,
+  "updated_at": "2026-05-13T06:05:39Z"
+}
\ No newline at end of file
diff --git a/wandb_uniform_transition_diag_history.jsonl b/wandb_uniform_transition_diag_history.jsonl
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391