diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..1ac16215676f7785ef78e19287115ac690bb53bb 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,10 @@ 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
+checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/09938061f7a5bc82b943e5a5b3de5611 filter=lfs diff=lfs merge=lfs -text
+checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/9e5368e04466de5ac9370801aead239f filter=lfs diff=lfs merge=lfs -text
+checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/fe0aed2012ead8e462701993e85ad88c filter=lfs diff=lfs merge=lfs -text
+checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/4736d5d81205aae60160cab8e4c4dd1f filter=lfs diff=lfs merge=lfs -text
+checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/a194174c60469a12768e95b81815ae4b filter=lfs diff=lfs merge=lfs -text
+checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/5c835200e1acc77bbeaaddbecbf1e692 filter=lfs diff=lfs merge=lfs -text
+checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/607c821719df130ce48bb086f0568ff8 filter=lfs diff=lfs merge=lfs -text
diff --git a/checkpoint_400000/checkpoint_400000/_CHECKPOINT_METADATA b/checkpoint_400000/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..ee9b59d997f9351f682c2942b90050fd19d73374
--- /dev/null
+++ b/checkpoint_400000/checkpoint_400000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778055618956314628, "commit_timestamp_nsecs": 1778055631347036247, "custom_metadata": {}}
\ No newline at end of file
diff --git a/checkpoint_400000/checkpoint_400000/_METADATA b/checkpoint_400000/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..1a316fda8d7906d1e8095306ad52f285a66a2744
--- /dev/null
+++ b/checkpoint_400000/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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024, 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": [2048]}}, "('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": [1024, 2048]}}, "('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, 1024]}}, "('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": [1024]}}, "('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, 1024]}}, "('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": [1024]}}, "('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, 1024]}}, "('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": [1024]}}, "('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": [1024, 1024]}}, "('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_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, 1024]}}, "('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": [24, 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": [25, 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": [1024, 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": [1024, 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": [1024]}}, "('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, 1024]}}}, "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/checkpoint_400000/_sharding b/checkpoint_400000/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..f2108a132bf94fbd7ca39fd30620a7be16596ce2
--- /dev/null
+++ b/checkpoint_400000/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))\"}","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))\"}","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/checkpoint_400000/array_metadatas/process_0 b/checkpoint_400000/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..20f836c2cd583c26da532110e5e4e75926261f00
--- /dev/null
+++ b/checkpoint_400000/checkpoint_400000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_12.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_13.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_14.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_15.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_16.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_17.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_18.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_19.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_20.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_21.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_22.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_23.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.DiTBlock_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "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": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.FinalLayer_0.Dense_1.kernel", "write_shape": [1024, 2048], "chunk_shape": [1024, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.LabelEmbedder_0.Embed_0.embedding", "write_shape": [1001, 1024], "chunk_shape": [1001, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.PatchedPatchEmbed_0.proj.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.PatchedPatchEmbed_0.proj.kernel", "write_shape": [16, 1024], "chunk_shape": [16, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_0.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_0.kernel", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "backbone.TimestepEmbedder_0.Dense_1.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "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_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, 1024], "chunk_shape": [1001, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_delta_embed.embedding", "write_shape": [24, 32], "chunk_shape": [24, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.cond_layer_embed.embedding", "write_shape": [25, 32], "chunk_shape": [25, 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": [1024, 32], "chunk_shape": [1024, 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": [1024, 480], "chunk_shape": [1024, 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": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "predictor.out_proj.kernel", "write_shape": [480, 1024], "chunk_shape": [480, 1024], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/checkpoint_400000/checkpoint_400000/d/5b3e40c15715269b9517fcc2afce0461 b/checkpoint_400000/checkpoint_400000/d/5b3e40c15715269b9517fcc2afce0461
new file mode 100644
index 0000000000000000000000000000000000000000..180e2ce4fb2cc7dba8052ec7840b50dbead55584
Binary files /dev/null and b/checkpoint_400000/checkpoint_400000/d/5b3e40c15715269b9517fcc2afce0461 differ
diff --git a/checkpoint_400000/checkpoint_400000/manifest.ocdbt b/checkpoint_400000/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..214195a1f3f330c4474a012c14c7c62763a60372
Binary files /dev/null and b/checkpoint_400000/checkpoint_400000/manifest.ocdbt differ
diff --git a/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/09938061f7a5bc82b943e5a5b3de5611 b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/09938061f7a5bc82b943e5a5b3de5611
new file mode 100644
index 0000000000000000000000000000000000000000..5cf4685b6b9caa8cad1c6a1dfb385fa00fe81a53
--- /dev/null
+++ b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/09938061f7a5bc82b943e5a5b3de5611
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63001137ab26bf9c7e2fd32a2672f20c89f6277a61150408e8c81054b79f842e
+size 322998272
diff --git a/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/3bf3557ed23ef74b95b8be048cdf9a6c b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/3bf3557ed23ef74b95b8be048cdf9a6c
new file mode 100644
index 0000000000000000000000000000000000000000..69ff2c9da95051510bb02e5bcbadd7c025f64eda
Binary files /dev/null and b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/3bf3557ed23ef74b95b8be048cdf9a6c differ
diff --git a/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/4d5c6ff5c3e0cbc34e52605dc684cc9e b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/4d5c6ff5c3e0cbc34e52605dc684cc9e
new file mode 100644
index 0000000000000000000000000000000000000000..b1c9d85a9ce283ba111ccfa2a2a5523a7c01f521
Binary files /dev/null and b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/4d5c6ff5c3e0cbc34e52605dc684cc9e differ
diff --git a/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/8f7a8b9568c0d403a57cdcba322cb6b6 b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/8f7a8b9568c0d403a57cdcba322cb6b6
new file mode 100644
index 0000000000000000000000000000000000000000..78d3ee874f5c7be24d4712c0e39fac0f9f603e0e
Binary files /dev/null and b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/8f7a8b9568c0d403a57cdcba322cb6b6 differ
diff --git a/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/9be05af07d48d789a36ba7e3f31aa9b4 b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/9be05af07d48d789a36ba7e3f31aa9b4
new file mode 100644
index 0000000000000000000000000000000000000000..0f8eca17a9198097eb0cd93c0b611de731f9b40a
Binary files /dev/null and b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/9be05af07d48d789a36ba7e3f31aa9b4 differ
diff --git a/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/9e5368e04466de5ac9370801aead239f b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/9e5368e04466de5ac9370801aead239f
new file mode 100644
index 0000000000000000000000000000000000000000..d269e634190f3cbad1dd4c98fe9d4d35749a9704
--- /dev/null
+++ b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/9e5368e04466de5ac9370801aead239f
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94f6d8913f7aa79d026d5c0d79d8683e9f96f808829903c02407db223a18de10
+size 50696192
diff --git a/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/f93a46d62bd9cbe88225f68855de4347 b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/f93a46d62bd9cbe88225f68855de4347
new file mode 100644
index 0000000000000000000000000000000000000000..8f2169eabed293bce57990ff3a39647f71653e4e
Binary files /dev/null and b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/f93a46d62bd9cbe88225f68855de4347 differ
diff --git a/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/fe0aed2012ead8e462701993e85ad88c b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/fe0aed2012ead8e462701993e85ad88c
new file mode 100644
index 0000000000000000000000000000000000000000..059bd0f70e9850957046ed44163d688755755c7d
--- /dev/null
+++ b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/d/fe0aed2012ead8e462701993e85ad88c
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fcfa38a01816ca139169e4b812c57a10f4f60a84d8c5ee317ac55ab2509129f9
+size 1429045248
diff --git a/checkpoint_400000/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..e6247edbf582379fbbc54506ece240642835e8fb
Binary files /dev/null and b/checkpoint_400000/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/checkpoint_400000/ema/checkpoint_400000/_CHECKPOINT_METADATA b/checkpoint_400000/ema/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..330c2a44a2c9ee329f5f90c2bd07e21b78751f95
--- /dev/null
+++ b/checkpoint_400000/ema/checkpoint_400000/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1778055654293161461, "commit_timestamp_nsecs": 1778055665461213791, "custom_metadata": {}}
\ No newline at end of file
diff --git a/checkpoint_400000/ema/checkpoint_400000/_METADATA b/checkpoint_400000/ema/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..36992cb7e46744ae9dc1f547f5d303520b7b1baf
--- /dev/null
+++ b/checkpoint_400000/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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [6144]}}, "('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": [1024, 6144]}}, "('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": [4096]}}, "('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": [1024, 4096]}}, "('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": [1024]}}, "('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": [4096, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024]}}, "('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, 64, 1024]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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, 64]}}, "('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": [1024, 16, 64]}}, "('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": [1024, 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": [2048]}}, "('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": [1024, 2048]}}, "('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, 1024]}}, "('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": [1024]}}, "('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, 1024]}}, "('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": [1024]}}, "('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, 1024]}}, "('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": [1024]}}, "('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": [1024, 1024]}}}, "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/ema/checkpoint_400000/_sharding b/checkpoint_400000/ema/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..c6295527515547b4ff433325689ac4cdc3aaa657
--- /dev/null
+++ b/checkpoint_400000/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))\"}","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/checkpoint_400000/ema/checkpoint_400000/array_metadatas/process_0 b/checkpoint_400000/ema/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..5597e96a37ad8deee5914fba28072dbe793be13b
--- /dev/null
+++ b/checkpoint_400000/ema/checkpoint_400000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "DiTBlock_0.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_0.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_1.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_10.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_11.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_12.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_13.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_14.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_15.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_16.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_17.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_18.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_19.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_2.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_20.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_21.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_22.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_23.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_3.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_4.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_5.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_6.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_7.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_8.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_0.bias", "write_shape": [6144], "chunk_shape": [6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_0.kernel", "write_shape": [1024, 6144], "chunk_shape": [1024, 6144], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_1.bias", "write_shape": [4096], "chunk_shape": [4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_1.kernel", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_2.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.Dense_2.kernel", "write_shape": [4096, 1024], "chunk_shape": [4096, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.key.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.key.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.out.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.out.kernel", "write_shape": [16, 64, 1024], "chunk_shape": [16, 64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.query.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.query.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.value.bias", "write_shape": [16, 64], "chunk_shape": [16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "DiTBlock_9.MultiHeadDotProductAttention_0.value.kernel", "write_shape": [1024, 16, 64], "chunk_shape": [1024, 16, 64], "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": [1024, 32], "chunk_shape": [1024, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "FinalLayer_0.Dense_1.bias", "write_shape": [2048], "chunk_shape": [2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "FinalLayer_0.Dense_1.kernel", "write_shape": [1024, 2048], "chunk_shape": [1024, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "LabelEmbedder_0.Embed_0.embedding", "write_shape": [1001, 1024], "chunk_shape": [1001, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "PatchedPatchEmbed_0.proj.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "PatchedPatchEmbed_0.proj.kernel", "write_shape": [16, 1024], "chunk_shape": [16, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "TimestepEmbedder_0.Dense_0.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "TimestepEmbedder_0.Dense_0.kernel", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "TimestepEmbedder_0.Dense_1.bias", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "TimestepEmbedder_0.Dense_1.kernel", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/checkpoint_400000/ema/checkpoint_400000/d/646bca6b0f4db1522edc62e9c4ddcdea b/checkpoint_400000/ema/checkpoint_400000/d/646bca6b0f4db1522edc62e9c4ddcdea
new file mode 100644
index 0000000000000000000000000000000000000000..cc9d4a1885877ad3540951360210ddceedf9b47b
Binary files /dev/null and b/checkpoint_400000/ema/checkpoint_400000/d/646bca6b0f4db1522edc62e9c4ddcdea differ
diff --git a/checkpoint_400000/ema/checkpoint_400000/manifest.ocdbt b/checkpoint_400000/ema/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..7e6f606096af5b527fda021ce7c37ec1a159787f
Binary files /dev/null and b/checkpoint_400000/ema/checkpoint_400000/manifest.ocdbt differ
diff --git a/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/0ec47fa62e6d5541e8b1d5acfa2eba7d b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/0ec47fa62e6d5541e8b1d5acfa2eba7d
new file mode 100644
index 0000000000000000000000000000000000000000..cff6b9c922762be97042abfb6db3eeb958232032
Binary files /dev/null and b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/0ec47fa62e6d5541e8b1d5acfa2eba7d differ
diff --git a/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/4736d5d81205aae60160cab8e4c4dd1f b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/4736d5d81205aae60160cab8e4c4dd1f
new file mode 100644
index 0000000000000000000000000000000000000000..e813a397369ac4f3db467672016d6056645c73a7
--- /dev/null
+++ b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/4736d5d81205aae60160cab8e4c4dd1f
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec32ae01eafe588bf4426eb7d11ba878d951caeaf11b6d5fc532a63c11d6986f
+size 823123968
diff --git a/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/899bcc292b412eb499defeaf69976075 b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/899bcc292b412eb499defeaf69976075
new file mode 100644
index 0000000000000000000000000000000000000000..94e56465234706e1a418c2bb3da6d70c413640c5
Binary files /dev/null and b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/899bcc292b412eb499defeaf69976075 differ
diff --git a/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/a194174c60469a12768e95b81815ae4b b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/a194174c60469a12768e95b81815ae4b
new file mode 100644
index 0000000000000000000000000000000000000000..cc44cdfbf0ccb3e60961ebe19464aa8f84ab2f27
--- /dev/null
+++ b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/a194174c60469a12768e95b81815ae4b
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e319bced6f052fec98a26be4eec997b4d780fe0d928d9d397ca183e523f82fc
+size 875900928
diff --git a/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/b7cf8e1e407842a0003d607df4bbcd5a b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/b7cf8e1e407842a0003d607df4bbcd5a
new file mode 100644
index 0000000000000000000000000000000000000000..d59c06744a94833be8df05980c668781af43c9b8
Binary files /dev/null and b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/b7cf8e1e407842a0003d607df4bbcd5a differ
diff --git a/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/e43729e96fcdea605ce3acfe0b9dd6d1 b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/e43729e96fcdea605ce3acfe0b9dd6d1
new file mode 100644
index 0000000000000000000000000000000000000000..1d89930ae16388fa6bf1a5725277ef965878e6aa
Binary files /dev/null and b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/d/e43729e96fcdea605ce3acfe0b9dd6d1 differ
diff --git a/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..ce0a8ad91c4a7867a58b3fd8dc6ee1c994333a14
Binary files /dev/null and b/checkpoint_400000/ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/checkpoint_400000/l2_ema/checkpoint_400000/_CHECKPOINT_METADATA b/checkpoint_400000/l2_ema/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..1ee0ec5938f9f5d4bb439323af0935b41768c038
--- /dev/null
+++ b/checkpoint_400000/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": 1778055666350075783, "commit_timestamp_nsecs": 1778055666369978903, "custom_metadata": {}}
\ No newline at end of file
diff --git a/checkpoint_400000/l2_ema/checkpoint_400000/_METADATA b/checkpoint_400000/l2_ema/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..d7c56100142fdc14791ac5804d35a817167d23c3
--- /dev/null
+++ b/checkpoint_400000/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": [25, 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/checkpoint_400000/l2_ema/checkpoint_400000/_sharding b/checkpoint_400000/l2_ema/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..ac31b218b77c613e621eacb64ab85a4c534b1351
--- /dev/null
+++ b/checkpoint_400000/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/checkpoint_400000/l2_ema/checkpoint_400000/array_metadatas/process_0 b/checkpoint_400000/l2_ema/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..c302954dcd2ec3dc6fcb73585d68165175b45d01
--- /dev/null
+++ b/checkpoint_400000/l2_ema/checkpoint_400000/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "l2_ema", "write_shape": [25, 5, 256], "chunk_shape": [25, 5, 256], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/checkpoint_400000/l2_ema/checkpoint_400000/d/16d9a81f1dd19b2c9180d8b06142852d b/checkpoint_400000/l2_ema/checkpoint_400000/d/16d9a81f1dd19b2c9180d8b06142852d
new file mode 100644
index 0000000000000000000000000000000000000000..03eef89e6eaadd8ba95e5c0ad0c8929864078c52
Binary files /dev/null and b/checkpoint_400000/l2_ema/checkpoint_400000/d/16d9a81f1dd19b2c9180d8b06142852d differ
diff --git a/checkpoint_400000/l2_ema/checkpoint_400000/manifest.ocdbt b/checkpoint_400000/l2_ema/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..805fdd5b984e87105389f59304b20167c8e5cea8
Binary files /dev/null and b/checkpoint_400000/l2_ema/checkpoint_400000/manifest.ocdbt differ
diff --git a/checkpoint_400000/l2_ema/checkpoint_400000/ocdbt.process_0/d/276a0247b90f8bdb3d7c985bc707f786 b/checkpoint_400000/l2_ema/checkpoint_400000/ocdbt.process_0/d/276a0247b90f8bdb3d7c985bc707f786
new file mode 100644
index 0000000000000000000000000000000000000000..03eef89e6eaadd8ba95e5c0ad0c8929864078c52
Binary files /dev/null and b/checkpoint_400000/l2_ema/checkpoint_400000/ocdbt.process_0/d/276a0247b90f8bdb3d7c985bc707f786 differ
diff --git a/checkpoint_400000/l2_ema/checkpoint_400000/ocdbt.process_0/d/ae466a4a13789a19318c2837eacc8774 b/checkpoint_400000/l2_ema/checkpoint_400000/ocdbt.process_0/d/ae466a4a13789a19318c2837eacc8774
new file mode 100644
index 0000000000000000000000000000000000000000..da1de4b52cf4083fc9b335e53bce1336834fbb98
Binary files /dev/null and b/checkpoint_400000/l2_ema/checkpoint_400000/ocdbt.process_0/d/ae466a4a13789a19318c2837eacc8774 differ
diff --git a/checkpoint_400000/l2_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/checkpoint_400000/l2_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..877c79a485c563876dd59284c4c784293515ef44
Binary files /dev/null and b/checkpoint_400000/l2_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/_CHECKPOINT_METADATA b/checkpoint_400000/predictor_ema/checkpoint_400000/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..0999462d8b1462bf284bb2856ec1fc2092c8a24d
--- /dev/null
+++ b/checkpoint_400000/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": 1778055665669374102, "commit_timestamp_nsecs": 1778055666332764263, "custom_metadata": {}}
\ No newline at end of file
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/_METADATA b/checkpoint_400000/predictor_ema/checkpoint_400000/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..2c02c99aa2bf383ee8d021dc42c7472cb280a4d7
--- /dev/null
+++ b/checkpoint_400000/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_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, 1024]}}, "('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": [24, 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": [25, 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": [1024, 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": [1024, 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": [1024]}}, "('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, 1024]}}}, "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/predictor_ema/checkpoint_400000/_sharding b/checkpoint_400000/predictor_ema/checkpoint_400000/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..25ec8a16328937a6ee6b1e81a371254531c2f54b
--- /dev/null
+++ b/checkpoint_400000/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))\"}","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/checkpoint_400000/predictor_ema/checkpoint_400000/array_metadatas/process_0 b/checkpoint_400000/predictor_ema/checkpoint_400000/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..b91f032e5831da1b0d36512567fb0d4988f3bbd7
--- /dev/null
+++ b/checkpoint_400000/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_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, 1024], "chunk_shape": [1001, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_delta_embed.embedding", "write_shape": [24, 32], "chunk_shape": [24, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "cond_layer_embed.embedding", "write_shape": [25, 32], "chunk_shape": [25, 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": [1024, 32], "chunk_shape": [1024, 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": [1024, 480], "chunk_shape": [1024, 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": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "out_proj.kernel", "write_shape": [480, 1024], "chunk_shape": [480, 1024], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/d/80d7c35d48b091700302b0f2a415442d b/checkpoint_400000/predictor_ema/checkpoint_400000/d/80d7c35d48b091700302b0f2a415442d
new file mode 100644
index 0000000000000000000000000000000000000000..84b13624b4d460b329f02f4400268839c568eef7
Binary files /dev/null and b/checkpoint_400000/predictor_ema/checkpoint_400000/d/80d7c35d48b091700302b0f2a415442d differ
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/manifest.ocdbt b/checkpoint_400000/predictor_ema/checkpoint_400000/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..4c40359e55f02b01db09eb9b2b0898cf75b11f23
Binary files /dev/null and b/checkpoint_400000/predictor_ema/checkpoint_400000/manifest.ocdbt differ
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/207f24ee9512364dca3dab1f881bde08 b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/207f24ee9512364dca3dab1f881bde08
new file mode 100644
index 0000000000000000000000000000000000000000..787f8a5c3de4e3e64f4c97d822e5ebc45f432eb0
Binary files /dev/null and b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/207f24ee9512364dca3dab1f881bde08 differ
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/5c835200e1acc77bbeaaddbecbf1e692 b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/5c835200e1acc77bbeaaddbecbf1e692
new file mode 100644
index 0000000000000000000000000000000000000000..1ccc7c4f87a97a6763ad79035d297eae333f8907
--- /dev/null
+++ b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/5c835200e1acc77bbeaaddbecbf1e692
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c54e22f750e92f1ff0c3d0b93c85233971ad79ab3793ede6d47c27f2da2c3cc3
+size 97312768
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/607c821719df130ce48bb086f0568ff8 b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/607c821719df130ce48bb086f0568ff8
new file mode 100644
index 0000000000000000000000000000000000000000..ef6c11597e06c0f5ddeb4448950c5da9418743a9
--- /dev/null
+++ b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/607c821719df130ce48bb086f0568ff8
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94d7dea8aa55c231d16ecf29e2228af2c8a97cf0b72d3d322d67f51ad1ca6df4
+size 6545408
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/6c4cb504eae533d9ecaf8aa5fe0ca07a b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/6c4cb504eae533d9ecaf8aa5fe0ca07a
new file mode 100644
index 0000000000000000000000000000000000000000..69262604cacb486cf7ae514dfb7c845e6229d79e
Binary files /dev/null and b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/6c4cb504eae533d9ecaf8aa5fe0ca07a differ
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/960a8764570ed39a06c52bb431d5f0c6 b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/960a8764570ed39a06c52bb431d5f0c6
new file mode 100644
index 0000000000000000000000000000000000000000..fb1498846434e79afe3482f241d103f3ddbb4528
Binary files /dev/null and b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/960a8764570ed39a06c52bb431d5f0c6 differ
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/aa595d364a2dd0e164e690870cf595dd b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/aa595d364a2dd0e164e690870cf595dd
new file mode 100644
index 0000000000000000000000000000000000000000..c7cfead1d92f187c3ecf07e86f0d975779f554a8
Binary files /dev/null and b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/aa595d364a2dd0e164e690870cf595dd differ
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/d420dfa421a97d4f8200a4e61145ce6a b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/d420dfa421a97d4f8200a4e61145ce6a
new file mode 100644
index 0000000000000000000000000000000000000000..e7383916619f695593d701291e0ff1f64c6a2a07
Binary files /dev/null and b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/d/d420dfa421a97d4f8200a4e61145ce6a differ
diff --git a/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..f237f51aeb277e0170673e6a36885966271aea8f
Binary files /dev/null and b/checkpoint_400000/predictor_ema/checkpoint_400000/ocdbt.process_0/manifest.ocdbt differ