diff --git a/.gitattributes b/.gitattributes
index cdf57a94747b40a2cf9fb0450ca72bf0c78b2ae9..2f119661e6fd54546a88740e5e9579dbf768311e 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -47,3 +47,8 @@ se-shadow_5000_40/ocdbt.process_0/d/a22ec15c03906f05880a97a239959288 filter=lfs
 se-shadow_5000_40/ocdbt.process_0/d/cea631c0961e02e6d7d8c6d41008b007 filter=lfs diff=lfs merge=lfs -text
 se-shadow_5000_40/ocdbt.process_0/d/e19441b2d810869b6aba275639b6e9f1 filter=lfs diff=lfs merge=lfs -text
 se-shadow_5000_40/ocdbt.process_0/d/f25af335e4995be9e41552388b05f4df filter=lfs diff=lfs merge=lfs -text
+me-full_25000_5/d/718edd575ca5903ee28e438c8edd7c6d filter=lfs diff=lfs merge=lfs -text
+me-full_25000_5/ocdbt.process_0/d/00af2b5dd8c97f8aa04c90314584796b filter=lfs diff=lfs merge=lfs -text
+me-full_25000_5/ocdbt.process_0/d/231e228c0807bec3e5e5dfa0fbe0deb1 filter=lfs diff=lfs merge=lfs -text
+me-full_25000_5/ocdbt.process_0/d/7f413e304642522d0dc6738b225998d7 filter=lfs diff=lfs merge=lfs -text
+me-full_25000_5/ocdbt.process_0/d/a150499bc1dda58bef676984f1c5c136 filter=lfs diff=lfs merge=lfs -text
diff --git a/me-full_25000_5/_CHECKPOINT_METADATA b/me-full_25000_5/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..78fa7194be95123899fd632a358c624baa39e2c8
--- /dev/null
+++ b/me-full_25000_5/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1783559112151463017, "commit_timestamp_nsecs": 1783559118540023382, "custom_metadata": {}}
\ No newline at end of file
diff --git a/me-full_25000_5/_METADATA b/me-full_25000_5/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..15bc840afb188b3f12a0cae4c24e9b252b0c1564
--- /dev/null
+++ b/me-full_25000_5/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('0', 'decoder', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'decoder', 'dof_gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'decoder', 'dof_gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'decoder', 'dof_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'decoder', 'dof_gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'decoder', 'dof_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'decoder', 'dof_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'decoder', 'ln_dist_dof', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_dist_dof', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_dist_pose', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_dist_pose', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_dof', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_dof', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_pose', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_pose', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'pos_head', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 7, 9]}}, "('0', 'decoder', 'pos_head', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 7, 9]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'pos_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('0', 'decoder', 'pos_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'decoder', 'pos_head', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'pos_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('0', 'decoder', 'pose_gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'decoder', 'pose_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 4]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'decoder', 'rot_head', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'rot_head', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 7, 9]}}, "('0', 'decoder', 'rot_head', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 7, 9]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'rot_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('0', 'decoder', 'rot_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'decoder', 'rot_head', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'rot_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'kinematic_encoding', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'kinematic_encoding', 'gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('0', 'kinematic_encoding', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'kinematic_encoding', 'gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'kinematic_encoding', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'kinematic_encoding', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'align_to_approach', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'align_to_approach', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [13, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'fingertip_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fingertip_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'fingertip_normals')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fingertip_normals", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'init_pregrasp_joint')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "init_pregrasp_joint", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'local_fingertip_contact_positions')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "local_fingertip_contact_positions", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 6, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'align_to_approach', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'align_to_approach', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [17, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'fingertip_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "fingertip_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'fingertip_normals')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "fingertip_normals", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'init_pregrasp_joint')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "init_pregrasp_joint", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'local_fingertip_contact_positions')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "local_fingertip_contact_positions", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 3, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [23, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [22, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [22, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [22, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'ln_dist', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_dist', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot_sh', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot_sh', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'kinematic_encoding', 'q_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'kinematic_encoding', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'kinematic_encoding', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_s_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_s_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'proj', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "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": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'proj', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'proj', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 32]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'proj', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'proj', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'proj', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'weight_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'weight_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'weight_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'weight_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'count')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [13, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [17, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [23, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot_sh', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot_sh', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'q_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "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": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [13, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [17, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [23, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot_sh', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot_sh', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'q_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "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, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "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": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "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": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '2', 'count')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('1', 'step')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "step", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}}, "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/me-full_25000_5/_sharding b/me-full_25000_5/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..74ab7ac219f1dc2740b738c2f142a64227d66d13
--- /dev/null
+++ b/me-full_25000_5/_sharding
@@ -0,0 +1 @@
+{"MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9pcnJlcC5ydG9p":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYWxpZ25fdG9fYXBwcm9hY2guMA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYWxpZ25fdG9fYXBwcm9hY2guMQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZmluZ2VydGlwX25vcm1hbHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZmluZ2VydGlwX2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuaW5pdF9wcmVncmFzcF9qb2ludA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIubG9jYWxfZmluZ2VydGlwX2NvbnRhY3RfcG9zaXRpb25z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYWxpZ25fdG9fYXBwcm9hY2guMA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYWxpZ25fdG9fYXBwcm9hY2guMQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZmluZ2VydGlwX25vcm1hbHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZmluZ2VydGlwX2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuaW5pdF9wcmVncmFzcF9qb2ludA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMubG9jYWxfZmluZ2VydGlwX2NvbnRhY3RfcG9zaXRpb25z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcud2VpZ2h0X2lycmVwLnJ0b2k=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5jb3VudA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMi5jb3VudA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5zdGVw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}"}
\ No newline at end of file
diff --git a/me-full_25000_5/array_metadatas/process_0 b/me-full_25000_5/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..904e2ec3f8cfab2426dff227f1565193905a1667
--- /dev/null
+++ b/me-full_25000_5/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "0.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.coefficient_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid.from_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid.to_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.map.to_m_var", "write_shape": [4, 4], "chunk_shape": [4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.coefficient_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid.from_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid.to_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.joint_ranges", "write_shape": [2, 2], "chunk_shape": [2, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.joint_transforms", "write_shape": [2, 6], "chunk_shape": [2, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.kinematics_transforms", "write_shape": [2, 7], "chunk_shape": [2, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.joint_ranges", "write_shape": [2, 2], "chunk_shape": [2, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.joint_transforms", "write_shape": [2, 6], "chunk_shape": [2, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.kinematics_transforms", "write_shape": [2, 7], "chunk_shape": [2, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.align_to_approach.0", "write_shape": [3, 3], "chunk_shape": [3, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.align_to_approach.1", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 9], "chunk_shape": [13, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.fingertip_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.fingertip_normals", "write_shape": [3, 3], "chunk_shape": [3, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.init_pregrasp_joint", "write_shape": [12], "chunk_shape": [12], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.joint_ranges", "write_shape": [12, 2], "chunk_shape": [12, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.joint_transforms", "write_shape": [12, 6], "chunk_shape": [12, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.kinematics_transforms", "write_shape": [12, 7], "chunk_shape": [12, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.local_fingertip_contact_positions", "write_shape": [3, 6, 3], "chunk_shape": [3, 6, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.align_to_approach.0", "write_shape": [3, 3], "chunk_shape": [3, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.align_to_approach.1", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 9], "chunk_shape": [17, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.fingertip_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.fingertip_normals", "write_shape": [4, 3], "chunk_shape": [4, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.init_pregrasp_joint", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.joint_ranges", "write_shape": [16, 2], "chunk_shape": [16, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.joint_transforms", "write_shape": [16, 6], "chunk_shape": [16, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.kinematics_transforms", "write_shape": [16, 7], "chunk_shape": [16, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.local_fingertip_contact_positions", "write_shape": [4, 3, 3], "chunk_shape": [4, 3, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 9], "chunk_shape": [23, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.joint_ranges", "write_shape": [22, 2], "chunk_shape": [22, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.joint_transforms", "write_shape": [22, 6], "chunk_shape": [22, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.kinematics_transforms", "write_shape": [22, 7], "chunk_shape": [22, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.joint_ranges", "write_shape": [1, 2], "chunk_shape": [1, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.joint_transforms", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.kinematics_transforms", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.expand_index", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.expand_index", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.weight", "write_shape": [3, 64, 32], "chunk_shape": [3, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.1.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.1.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 9], "chunk_shape": [13, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 9], "chunk_shape": [17, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 9], "chunk_shape": [23, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.proj.weight", "write_shape": [3, 64, 32], "chunk_shape": [3, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.proj.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.1.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.1.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 9], "chunk_shape": [13, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 9], "chunk_shape": [17, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 9], "chunk_shape": [23, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.proj.weight", "write_shape": [3, 64, 32], "chunk_shape": [3, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.proj.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.2.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/me-full_25000_5/d/718edd575ca5903ee28e438c8edd7c6d b/me-full_25000_5/d/718edd575ca5903ee28e438c8edd7c6d
new file mode 100644
index 0000000000000000000000000000000000000000..657516476eeff9f9491ce2f63ba83dcd9e3f769f
--- /dev/null
+++ b/me-full_25000_5/d/718edd575ca5903ee28e438c8edd7c6d
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d4697939b11f82494f8460ad0dd65e1cfd44c2afc7c594732beb1727a2e7c37
+size 718624
diff --git a/me-full_25000_5/manifest.ocdbt b/me-full_25000_5/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..46cb42ad86c4aa328e1ce509f432d55fd73be19f
Binary files /dev/null and b/me-full_25000_5/manifest.ocdbt differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/00af2b5dd8c97f8aa04c90314584796b b/me-full_25000_5/ocdbt.process_0/d/00af2b5dd8c97f8aa04c90314584796b
new file mode 100644
index 0000000000000000000000000000000000000000..948d3d149a0ebf4ab4f3cd67856911065e9abc35
--- /dev/null
+++ b/me-full_25000_5/ocdbt.process_0/d/00af2b5dd8c97f8aa04c90314584796b
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b17f2f885b9a0570acb84b9683d3887f2651164df4a5a32b9c9038259d9f297
+size 748878
diff --git a/me-full_25000_5/ocdbt.process_0/d/0349ce9f86f0bd7a9aedbef8203be4ba b/me-full_25000_5/ocdbt.process_0/d/0349ce9f86f0bd7a9aedbef8203be4ba
new file mode 100644
index 0000000000000000000000000000000000000000..30a35b081106bd9679805ec280e48b2dc577c3f7
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/0349ce9f86f0bd7a9aedbef8203be4ba differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/03dd430ba8369c48a87fcb03e42c7eac b/me-full_25000_5/ocdbt.process_0/d/03dd430ba8369c48a87fcb03e42c7eac
new file mode 100644
index 0000000000000000000000000000000000000000..9288f8288d573d2924e92939d831e85382281718
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/03dd430ba8369c48a87fcb03e42c7eac differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/0a8628e9a5ad66dddfe056d8c885b791 b/me-full_25000_5/ocdbt.process_0/d/0a8628e9a5ad66dddfe056d8c885b791
new file mode 100644
index 0000000000000000000000000000000000000000..0237521fdf46611ba7850ee7f57a9efee6dc5453
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/0a8628e9a5ad66dddfe056d8c885b791 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/0e37712ee178723dd0e4e2acfca85627 b/me-full_25000_5/ocdbt.process_0/d/0e37712ee178723dd0e4e2acfca85627
new file mode 100644
index 0000000000000000000000000000000000000000..4752a15f612ded6d4febbe50ef62ec8ee741c03c
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/0e37712ee178723dd0e4e2acfca85627 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/14ef2506af566cb93b0be19675ec3119 b/me-full_25000_5/ocdbt.process_0/d/14ef2506af566cb93b0be19675ec3119
new file mode 100644
index 0000000000000000000000000000000000000000..317b1dd421c9d3747c1aaa04736995fea8bdc7c8
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/14ef2506af566cb93b0be19675ec3119 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/18464eb53da14fb1f4bcc7f172174f50 b/me-full_25000_5/ocdbt.process_0/d/18464eb53da14fb1f4bcc7f172174f50
new file mode 100644
index 0000000000000000000000000000000000000000..766f51fc6e21c5e14ab03dcc032cc281b790a4df
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/18464eb53da14fb1f4bcc7f172174f50 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/1b15fbed51fcf73c2b8c58bdf2902add b/me-full_25000_5/ocdbt.process_0/d/1b15fbed51fcf73c2b8c58bdf2902add
new file mode 100644
index 0000000000000000000000000000000000000000..3240d6ed726877c798cb2cf2234b4542a0648afd
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/1b15fbed51fcf73c2b8c58bdf2902add differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/1b7fd50fa992086b54b19a26df85b5ba b/me-full_25000_5/ocdbt.process_0/d/1b7fd50fa992086b54b19a26df85b5ba
new file mode 100644
index 0000000000000000000000000000000000000000..f14e77a1a4098a03f3982314d8ce65085588d84e
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/1b7fd50fa992086b54b19a26df85b5ba differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/231e228c0807bec3e5e5dfa0fbe0deb1 b/me-full_25000_5/ocdbt.process_0/d/231e228c0807bec3e5e5dfa0fbe0deb1
new file mode 100644
index 0000000000000000000000000000000000000000..3ca5767d5d416ba6ded93cccc907ebf0974938e4
--- /dev/null
+++ b/me-full_25000_5/ocdbt.process_0/d/231e228c0807bec3e5e5dfa0fbe0deb1
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b227eda117e23c76b0ce75c69a60f25a0482694e558d2d3a835fe865b833886
+size 3803540
diff --git a/me-full_25000_5/ocdbt.process_0/d/296ed80ad39551f938ec242eb7eba2aa b/me-full_25000_5/ocdbt.process_0/d/296ed80ad39551f938ec242eb7eba2aa
new file mode 100644
index 0000000000000000000000000000000000000000..be082cc167ac5c96a525eb4028abae7dcdf557ee
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/296ed80ad39551f938ec242eb7eba2aa differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/29c40bfad25776ae7b50ba8de59cd163 b/me-full_25000_5/ocdbt.process_0/d/29c40bfad25776ae7b50ba8de59cd163
new file mode 100644
index 0000000000000000000000000000000000000000..15ca105a11cf4b2d357caf4d5bcf5a4de17c9509
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/29c40bfad25776ae7b50ba8de59cd163 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/2bd96fb2e186ead94b38181dcfbf43f0 b/me-full_25000_5/ocdbt.process_0/d/2bd96fb2e186ead94b38181dcfbf43f0
new file mode 100644
index 0000000000000000000000000000000000000000..1ea23944b50a554039e18deaa67a9d1a210e8622
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/2bd96fb2e186ead94b38181dcfbf43f0 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/2d19e40749eb94b8beb796c96930aa99 b/me-full_25000_5/ocdbt.process_0/d/2d19e40749eb94b8beb796c96930aa99
new file mode 100644
index 0000000000000000000000000000000000000000..8cabc290c1810148236e91691ab40c159f494887
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/2d19e40749eb94b8beb796c96930aa99 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/2f4d4f5b9e420e205d1e9dcb095a8b69 b/me-full_25000_5/ocdbt.process_0/d/2f4d4f5b9e420e205d1e9dcb095a8b69
new file mode 100644
index 0000000000000000000000000000000000000000..32966d8cf94254d5ffea9f763276dc88c5c1cf64
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/2f4d4f5b9e420e205d1e9dcb095a8b69 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/357e98659a44d93227ca8087ad516cd4 b/me-full_25000_5/ocdbt.process_0/d/357e98659a44d93227ca8087ad516cd4
new file mode 100644
index 0000000000000000000000000000000000000000..a4d734a58c0b4fbedfca07bbee8d8162f4f6d6e0
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/357e98659a44d93227ca8087ad516cd4 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/35ddf88b26e2fad2b53111b93c20f1ac b/me-full_25000_5/ocdbt.process_0/d/35ddf88b26e2fad2b53111b93c20f1ac
new file mode 100644
index 0000000000000000000000000000000000000000..04382754bd162d39d6707ec5ff44cdfc4c660072
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/35ddf88b26e2fad2b53111b93c20f1ac differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/36920e66aca9baafa9be692bc5a5b1a1 b/me-full_25000_5/ocdbt.process_0/d/36920e66aca9baafa9be692bc5a5b1a1
new file mode 100644
index 0000000000000000000000000000000000000000..04b1dd1c977f0bb410cdddac852eba930f7f19c3
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/36920e66aca9baafa9be692bc5a5b1a1 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/3aa3000b58633592a3ef8868e84a1cb5 b/me-full_25000_5/ocdbt.process_0/d/3aa3000b58633592a3ef8868e84a1cb5
new file mode 100644
index 0000000000000000000000000000000000000000..d852ba7b3270f474ff5a33cb1fd668851c0244bb
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/3aa3000b58633592a3ef8868e84a1cb5 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/3afd3f79fe68ac5255257caf9b4c1d27 b/me-full_25000_5/ocdbt.process_0/d/3afd3f79fe68ac5255257caf9b4c1d27
new file mode 100644
index 0000000000000000000000000000000000000000..8a58aad4d2c4d53c217ba63156d71b84fea45e70
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/3afd3f79fe68ac5255257caf9b4c1d27 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/3d2444b219204efd95eb776dba01fbfb b/me-full_25000_5/ocdbt.process_0/d/3d2444b219204efd95eb776dba01fbfb
new file mode 100644
index 0000000000000000000000000000000000000000..31f0edbf1e0f33e2073596032f84b43b5eb4b7fa
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/3d2444b219204efd95eb776dba01fbfb differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/3d9b39ae0c7562638ec5737df4a5481d b/me-full_25000_5/ocdbt.process_0/d/3d9b39ae0c7562638ec5737df4a5481d
new file mode 100644
index 0000000000000000000000000000000000000000..dc55e06752de4151fcb43a356e40d7f8b08991da
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/3d9b39ae0c7562638ec5737df4a5481d differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/420ba779ce2fe01b91e26be4f0a6bc97 b/me-full_25000_5/ocdbt.process_0/d/420ba779ce2fe01b91e26be4f0a6bc97
new file mode 100644
index 0000000000000000000000000000000000000000..690beb2c1e489c9245557fc8dbeffa5513e4283a
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/420ba779ce2fe01b91e26be4f0a6bc97 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/45c94aab6910b518761939c85a59d2f3 b/me-full_25000_5/ocdbt.process_0/d/45c94aab6910b518761939c85a59d2f3
new file mode 100644
index 0000000000000000000000000000000000000000..6582829aaabaf999c4ad7e49d79a289a18ad88b2
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/45c94aab6910b518761939c85a59d2f3 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/45e7dfe1e348c9e13c5f14203deb172b b/me-full_25000_5/ocdbt.process_0/d/45e7dfe1e348c9e13c5f14203deb172b
new file mode 100644
index 0000000000000000000000000000000000000000..c676462bb9bf224dbbec7a4efe0a968e0c1ec5f2
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/45e7dfe1e348c9e13c5f14203deb172b differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/48a2f2f4c053f0a2e2216eb20eaef88d b/me-full_25000_5/ocdbt.process_0/d/48a2f2f4c053f0a2e2216eb20eaef88d
new file mode 100644
index 0000000000000000000000000000000000000000..c11b2c85f93e908e1e94ebada12891b229411c2a
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/48a2f2f4c053f0a2e2216eb20eaef88d differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/51e90b15849208a5f745541c0dffedf1 b/me-full_25000_5/ocdbt.process_0/d/51e90b15849208a5f745541c0dffedf1
new file mode 100644
index 0000000000000000000000000000000000000000..324a80a7f08ad05d79566d90658581aa831a290b
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/51e90b15849208a5f745541c0dffedf1 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/58da2ceee3abafdb10b0b91d45240e50 b/me-full_25000_5/ocdbt.process_0/d/58da2ceee3abafdb10b0b91d45240e50
new file mode 100644
index 0000000000000000000000000000000000000000..db807276218e3bf7eb98f5a4f9fc2e1b65f0a50f
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/58da2ceee3abafdb10b0b91d45240e50 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/58f1fb3c61ba12abaa0c6d05296c400f b/me-full_25000_5/ocdbt.process_0/d/58f1fb3c61ba12abaa0c6d05296c400f
new file mode 100644
index 0000000000000000000000000000000000000000..0fd96d16d3393fbbd9cd22db90ebe947b92f6e36
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/58f1fb3c61ba12abaa0c6d05296c400f differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/61442ee479e06aae10b60e0e5e84f1ee b/me-full_25000_5/ocdbt.process_0/d/61442ee479e06aae10b60e0e5e84f1ee
new file mode 100644
index 0000000000000000000000000000000000000000..8800d133cdc2fd5c9958941e21a21891f59ae86c
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/61442ee479e06aae10b60e0e5e84f1ee differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/6206bfd2c4a5f0cc2b3744a75cfac537 b/me-full_25000_5/ocdbt.process_0/d/6206bfd2c4a5f0cc2b3744a75cfac537
new file mode 100644
index 0000000000000000000000000000000000000000..5da47573d450c62d481a0590a4033355137ab113
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/6206bfd2c4a5f0cc2b3744a75cfac537 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/67fd96db74c54e50a089d223623e37ed b/me-full_25000_5/ocdbt.process_0/d/67fd96db74c54e50a089d223623e37ed
new file mode 100644
index 0000000000000000000000000000000000000000..9b92284560b8328782795d91a21189dfd6b7a9c1
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/67fd96db74c54e50a089d223623e37ed differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/69374a1ffea7745ceea05d66867ef052 b/me-full_25000_5/ocdbt.process_0/d/69374a1ffea7745ceea05d66867ef052
new file mode 100644
index 0000000000000000000000000000000000000000..afe3ed94a8fe041b74ba96b289ebaf10e5e1366a
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/69374a1ffea7745ceea05d66867ef052 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/6a7a21f99db467d39d3d5929575b150a b/me-full_25000_5/ocdbt.process_0/d/6a7a21f99db467d39d3d5929575b150a
new file mode 100644
index 0000000000000000000000000000000000000000..788fbe6bb721914c1ae117b9ad173ce859d68bf1
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/6a7a21f99db467d39d3d5929575b150a differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/6a7c42c96efaed0b34141b5844b0202a b/me-full_25000_5/ocdbt.process_0/d/6a7c42c96efaed0b34141b5844b0202a
new file mode 100644
index 0000000000000000000000000000000000000000..e2acb87fb078c558f42ce11664f4e6c799d85434
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/6a7c42c96efaed0b34141b5844b0202a differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/6aec694a46199eaf65485fb360074860 b/me-full_25000_5/ocdbt.process_0/d/6aec694a46199eaf65485fb360074860
new file mode 100644
index 0000000000000000000000000000000000000000..fe37e8c663461ef68ca888f812476cb1844ed496
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/6aec694a46199eaf65485fb360074860 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/6dc243e1b7cdff04254f4082ffe6eb77 b/me-full_25000_5/ocdbt.process_0/d/6dc243e1b7cdff04254f4082ffe6eb77
new file mode 100644
index 0000000000000000000000000000000000000000..732b9c9bde4329e28616bb14cdde6fd5866cfd24
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/6dc243e1b7cdff04254f4082ffe6eb77 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/6f893e44ba2cdf3b070113d972fb0fe5 b/me-full_25000_5/ocdbt.process_0/d/6f893e44ba2cdf3b070113d972fb0fe5
new file mode 100644
index 0000000000000000000000000000000000000000..c15fc243a8d5147f553a82824bb8a64aa4e1e855
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/6f893e44ba2cdf3b070113d972fb0fe5 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/770ba564c07993f37dc98abac96b1b55 b/me-full_25000_5/ocdbt.process_0/d/770ba564c07993f37dc98abac96b1b55
new file mode 100644
index 0000000000000000000000000000000000000000..2bf3e3a6a8db0699accf80faf6515c0d7a4c711f
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/770ba564c07993f37dc98abac96b1b55 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/783c242e0adf249fe0308fa82275c7e5 b/me-full_25000_5/ocdbt.process_0/d/783c242e0adf249fe0308fa82275c7e5
new file mode 100644
index 0000000000000000000000000000000000000000..60c7ca5e543e8f1e21368fc78a388f8d8bf4868d
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/783c242e0adf249fe0308fa82275c7e5 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/7aadd0ee5ddc289a9c94ccbd4a3df906 b/me-full_25000_5/ocdbt.process_0/d/7aadd0ee5ddc289a9c94ccbd4a3df906
new file mode 100644
index 0000000000000000000000000000000000000000..d6c3bc18d8aac7c0d33020edc000db515060c2fe
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/7aadd0ee5ddc289a9c94ccbd4a3df906 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/7f413e304642522d0dc6738b225998d7 b/me-full_25000_5/ocdbt.process_0/d/7f413e304642522d0dc6738b225998d7
new file mode 100644
index 0000000000000000000000000000000000000000..ed156609ad8c184a96c825f5814ce1623efd514a
--- /dev/null
+++ b/me-full_25000_5/ocdbt.process_0/d/7f413e304642522d0dc6738b225998d7
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fcd79efa1bb7bf5a7591476d26c69b0220c5b7f4b92edd6ea739b5cabe3010ba
+size 114069504
diff --git a/me-full_25000_5/ocdbt.process_0/d/80fa9171f9ce3ab946ace596aa44b276 b/me-full_25000_5/ocdbt.process_0/d/80fa9171f9ce3ab946ace596aa44b276
new file mode 100644
index 0000000000000000000000000000000000000000..febafafb4389a8cba0cca22d269cea900289c266
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/80fa9171f9ce3ab946ace596aa44b276 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/8830ba4f6de3b21cc398adf5e215d461 b/me-full_25000_5/ocdbt.process_0/d/8830ba4f6de3b21cc398adf5e215d461
new file mode 100644
index 0000000000000000000000000000000000000000..67f714c7d29283c56af3cae5358d1234e45f6872
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/8830ba4f6de3b21cc398adf5e215d461 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/8e67b53c59d5a5c9f67305f523c51d52 b/me-full_25000_5/ocdbt.process_0/d/8e67b53c59d5a5c9f67305f523c51d52
new file mode 100644
index 0000000000000000000000000000000000000000..39d049cb223abb248fd5c11252ed4aed3879e0b9
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/8e67b53c59d5a5c9f67305f523c51d52 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/8f34f7140e885c0269ed110a5cc29b8f b/me-full_25000_5/ocdbt.process_0/d/8f34f7140e885c0269ed110a5cc29b8f
new file mode 100644
index 0000000000000000000000000000000000000000..7665b4aa372ed101a680baa445a0d41be416601f
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/8f34f7140e885c0269ed110a5cc29b8f differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/8fbdcf184b0d9e916713c3695c1a4337 b/me-full_25000_5/ocdbt.process_0/d/8fbdcf184b0d9e916713c3695c1a4337
new file mode 100644
index 0000000000000000000000000000000000000000..41bccb312c5f4fee7e193b1d6635c46dd1fec5d9
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/8fbdcf184b0d9e916713c3695c1a4337 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/969582610f05f7b728c81008d83ed709 b/me-full_25000_5/ocdbt.process_0/d/969582610f05f7b728c81008d83ed709
new file mode 100644
index 0000000000000000000000000000000000000000..0b1cecd28f72c1ebfb120658db03088e5538113c
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/969582610f05f7b728c81008d83ed709 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/96f12187290da43384ecc71e2597760e b/me-full_25000_5/ocdbt.process_0/d/96f12187290da43384ecc71e2597760e
new file mode 100644
index 0000000000000000000000000000000000000000..611abc6a0c574c3cd1e971a958ab636c5ba9df04
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/96f12187290da43384ecc71e2597760e differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/97b42a9c5573f683c5c4686f9daacad2 b/me-full_25000_5/ocdbt.process_0/d/97b42a9c5573f683c5c4686f9daacad2
new file mode 100644
index 0000000000000000000000000000000000000000..430314f851bec33c483eedbfdc2e791d769bd9c1
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/97b42a9c5573f683c5c4686f9daacad2 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/9bb3c5c3674676c4605ed3af3369020c b/me-full_25000_5/ocdbt.process_0/d/9bb3c5c3674676c4605ed3af3369020c
new file mode 100644
index 0000000000000000000000000000000000000000..0013f797fafe390e4713e63cdf787810412d7a93
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/9bb3c5c3674676c4605ed3af3369020c differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/9beadfd632931a4bfb5b287db46dffa7 b/me-full_25000_5/ocdbt.process_0/d/9beadfd632931a4bfb5b287db46dffa7
new file mode 100644
index 0000000000000000000000000000000000000000..a1b0b42511bfd58165f6c34dcedb9f725f6ebe01
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/9beadfd632931a4bfb5b287db46dffa7 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/9c5256779042461bc8ae29bea059050a b/me-full_25000_5/ocdbt.process_0/d/9c5256779042461bc8ae29bea059050a
new file mode 100644
index 0000000000000000000000000000000000000000..d561db6a4dbb7e99c8c9056649d2e7d6b0c9c472
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/9c5256779042461bc8ae29bea059050a differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/9dd493625fb804e3303428084d9dc17a b/me-full_25000_5/ocdbt.process_0/d/9dd493625fb804e3303428084d9dc17a
new file mode 100644
index 0000000000000000000000000000000000000000..e4aa42cbc1edbac31ca1ca6571f463ee82241d29
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/9dd493625fb804e3303428084d9dc17a differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/a150499bc1dda58bef676984f1c5c136 b/me-full_25000_5/ocdbt.process_0/d/a150499bc1dda58bef676984f1c5c136
new file mode 100644
index 0000000000000000000000000000000000000000..f204973ab36faa1df7d5e263c2f82293061f24b8
--- /dev/null
+++ b/me-full_25000_5/ocdbt.process_0/d/a150499bc1dda58bef676984f1c5c136
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0fb9ab885a3d8f2ee2df4c250493eb6dbd54b0f5d8c1b5875ecc0a74c24a8810
+size 89288704
diff --git a/me-full_25000_5/ocdbt.process_0/d/a65699b081ceb83970bcb588751361cc b/me-full_25000_5/ocdbt.process_0/d/a65699b081ceb83970bcb588751361cc
new file mode 100644
index 0000000000000000000000000000000000000000..d36fa77bce14bf6634defb777a0b147a1decc3d3
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/a65699b081ceb83970bcb588751361cc differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/a6b1a242064b4021b594a948d1c95bbb b/me-full_25000_5/ocdbt.process_0/d/a6b1a242064b4021b594a948d1c95bbb
new file mode 100644
index 0000000000000000000000000000000000000000..cc057c4eedf1556cd1b224cc0e78a0529cbc5682
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/a6b1a242064b4021b594a948d1c95bbb differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/a7081a867ca6c2c47afe6e25082f9702 b/me-full_25000_5/ocdbt.process_0/d/a7081a867ca6c2c47afe6e25082f9702
new file mode 100644
index 0000000000000000000000000000000000000000..91a4bc5094b59eac7b34a40972097aa9cc002d98
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/a7081a867ca6c2c47afe6e25082f9702 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/af16a7b863f586adb9bddb8e093ef61e b/me-full_25000_5/ocdbt.process_0/d/af16a7b863f586adb9bddb8e093ef61e
new file mode 100644
index 0000000000000000000000000000000000000000..cbb1b270c4ddce37fb1f09320bf64c15e7dc31fd
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/af16a7b863f586adb9bddb8e093ef61e differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/af5948fcfedb0a6085c6d4173b3e1925 b/me-full_25000_5/ocdbt.process_0/d/af5948fcfedb0a6085c6d4173b3e1925
new file mode 100644
index 0000000000000000000000000000000000000000..0d7e0de3ebfdaf6bda4d9bcbeb03d6eb15ad4f33
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/af5948fcfedb0a6085c6d4173b3e1925 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/b14d2b6b820213b9b18ea316b3bb575f b/me-full_25000_5/ocdbt.process_0/d/b14d2b6b820213b9b18ea316b3bb575f
new file mode 100644
index 0000000000000000000000000000000000000000..1bd0c0c9c2a85086f1d8e57a6b409bd919775095
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/b14d2b6b820213b9b18ea316b3bb575f differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/b1a50c0233fe89f6c550b08a857ad1e4 b/me-full_25000_5/ocdbt.process_0/d/b1a50c0233fe89f6c550b08a857ad1e4
new file mode 100644
index 0000000000000000000000000000000000000000..de07e8472afa384dab547efaf84aac40603921d1
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/b1a50c0233fe89f6c550b08a857ad1e4 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/b45e792f129d77ae43d197fe265fa22e b/me-full_25000_5/ocdbt.process_0/d/b45e792f129d77ae43d197fe265fa22e
new file mode 100644
index 0000000000000000000000000000000000000000..ae94c21615eb420334e294f6bc770ffba93ce7de
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/b45e792f129d77ae43d197fe265fa22e differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/b7ce07c5a33ef92243d5c5c1b02cc646 b/me-full_25000_5/ocdbt.process_0/d/b7ce07c5a33ef92243d5c5c1b02cc646
new file mode 100644
index 0000000000000000000000000000000000000000..463a365fb96e34995eaaca4197f7798fdb4eb61e
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/b7ce07c5a33ef92243d5c5c1b02cc646 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/b9d98839ab427f0c2b78ff8410af3ba3 b/me-full_25000_5/ocdbt.process_0/d/b9d98839ab427f0c2b78ff8410af3ba3
new file mode 100644
index 0000000000000000000000000000000000000000..e2f65b6aa1eca89cfa3313c461c14f01cec41b7b
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/b9d98839ab427f0c2b78ff8410af3ba3 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/bbaf678966f84beb8d4d94339d388706 b/me-full_25000_5/ocdbt.process_0/d/bbaf678966f84beb8d4d94339d388706
new file mode 100644
index 0000000000000000000000000000000000000000..cb1c508a5ab46cdcca7263f3a491e8d40145f0b6
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/bbaf678966f84beb8d4d94339d388706 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/c063456567bcfb0e2105029d6ce3b40f b/me-full_25000_5/ocdbt.process_0/d/c063456567bcfb0e2105029d6ce3b40f
new file mode 100644
index 0000000000000000000000000000000000000000..f9be7e9fef13c5850e79b437209e45733f7d7afd
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/c063456567bcfb0e2105029d6ce3b40f differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/c76c163d7ea600f79b2cd2df2e9125ae b/me-full_25000_5/ocdbt.process_0/d/c76c163d7ea600f79b2cd2df2e9125ae
new file mode 100644
index 0000000000000000000000000000000000000000..7377331eb85634f40d2825fddab4655d70f70e0e
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/c76c163d7ea600f79b2cd2df2e9125ae differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/c855e12d41641ded4ba71e4e0fd2b60d b/me-full_25000_5/ocdbt.process_0/d/c855e12d41641ded4ba71e4e0fd2b60d
new file mode 100644
index 0000000000000000000000000000000000000000..6bb2ca20dff70e8dc1f13e22ed83c15d1a64bff3
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/c855e12d41641ded4ba71e4e0fd2b60d differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/cb0381a4a8a09675d0e4162aec43f828 b/me-full_25000_5/ocdbt.process_0/d/cb0381a4a8a09675d0e4162aec43f828
new file mode 100644
index 0000000000000000000000000000000000000000..cdd3038840c44bf5d9bf00f7863ecc1ad8f02392
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/cb0381a4a8a09675d0e4162aec43f828 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/cd7ae19eb329d40d497e0aae28085760 b/me-full_25000_5/ocdbt.process_0/d/cd7ae19eb329d40d497e0aae28085760
new file mode 100644
index 0000000000000000000000000000000000000000..79c90666e0a5c81da9ccebdb33236a2f075477c9
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/cd7ae19eb329d40d497e0aae28085760 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/ce97faa2089d44d8e7725ce4162b7ce9 b/me-full_25000_5/ocdbt.process_0/d/ce97faa2089d44d8e7725ce4162b7ce9
new file mode 100644
index 0000000000000000000000000000000000000000..97d0598bff64609836230fe6c1b8936e9c3cc706
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/ce97faa2089d44d8e7725ce4162b7ce9 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/d1f0611d6981272d17eb9cee28d8a48d b/me-full_25000_5/ocdbt.process_0/d/d1f0611d6981272d17eb9cee28d8a48d
new file mode 100644
index 0000000000000000000000000000000000000000..9c7bc5f4830888e958c2aea4404044b14d0aa684
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/d1f0611d6981272d17eb9cee28d8a48d differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/d5527a0c41c7f685a43b671476aa68c6 b/me-full_25000_5/ocdbt.process_0/d/d5527a0c41c7f685a43b671476aa68c6
new file mode 100644
index 0000000000000000000000000000000000000000..6710d44d6a80d071f2c3cfcfba7ab51426a4d997
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/d5527a0c41c7f685a43b671476aa68c6 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/d8115d9803d64adcf0fddfb6cd5130bb b/me-full_25000_5/ocdbt.process_0/d/d8115d9803d64adcf0fddfb6cd5130bb
new file mode 100644
index 0000000000000000000000000000000000000000..4e86f2508b6e3da28eee613ba639ec0d58d03587
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/d8115d9803d64adcf0fddfb6cd5130bb differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/d820b21b6e3c6949958b463b21ca05c5 b/me-full_25000_5/ocdbt.process_0/d/d820b21b6e3c6949958b463b21ca05c5
new file mode 100644
index 0000000000000000000000000000000000000000..48c6011db02e435051542b6889e987472ffbf0db
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/d820b21b6e3c6949958b463b21ca05c5 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/dc0d5b857c97ef28be19646f8ba1049e b/me-full_25000_5/ocdbt.process_0/d/dc0d5b857c97ef28be19646f8ba1049e
new file mode 100644
index 0000000000000000000000000000000000000000..7a7c2449d7348d6695d8943838e872bbdfb94063
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/dc0d5b857c97ef28be19646f8ba1049e differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/dc588ae5756b303e358398db13b08ac5 b/me-full_25000_5/ocdbt.process_0/d/dc588ae5756b303e358398db13b08ac5
new file mode 100644
index 0000000000000000000000000000000000000000..ba38333efbebab8186b709b7d87c3a93c103a057
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/dc588ae5756b303e358398db13b08ac5 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/dc7fec401d8386ec6dc2a8da5a6db4b2 b/me-full_25000_5/ocdbt.process_0/d/dc7fec401d8386ec6dc2a8da5a6db4b2
new file mode 100644
index 0000000000000000000000000000000000000000..3fc1e10628bdbde536b1af5b7a4a1d68d392d1b7
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/dc7fec401d8386ec6dc2a8da5a6db4b2 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/df2f90faa806acfb2bedaae9443d2b8b b/me-full_25000_5/ocdbt.process_0/d/df2f90faa806acfb2bedaae9443d2b8b
new file mode 100644
index 0000000000000000000000000000000000000000..f84a6551a3461435a3712baab78b1a2c2b8f2656
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/df2f90faa806acfb2bedaae9443d2b8b differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/e286f5520ce45eb641464bb2b4ebcb8e b/me-full_25000_5/ocdbt.process_0/d/e286f5520ce45eb641464bb2b4ebcb8e
new file mode 100644
index 0000000000000000000000000000000000000000..d1af40abdadc0280b5265d7886e9bf871e8e0a99
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/e286f5520ce45eb641464bb2b4ebcb8e differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/e37f88c887972c45f8931803cdfeec88 b/me-full_25000_5/ocdbt.process_0/d/e37f88c887972c45f8931803cdfeec88
new file mode 100644
index 0000000000000000000000000000000000000000..e9ba169cb24b01a898f2d71030160a2753b026e9
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/e37f88c887972c45f8931803cdfeec88 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/e71df03f5ee4a6b5525c97fd58530cac b/me-full_25000_5/ocdbt.process_0/d/e71df03f5ee4a6b5525c97fd58530cac
new file mode 100644
index 0000000000000000000000000000000000000000..18526b74a22ca667957e1e196ad41cea5a54eaf9
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/e71df03f5ee4a6b5525c97fd58530cac differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/e7bd4f139529e30fb5984d58ee5a942c b/me-full_25000_5/ocdbt.process_0/d/e7bd4f139529e30fb5984d58ee5a942c
new file mode 100644
index 0000000000000000000000000000000000000000..e1c14b619c9920a4158c3c8de2ebdac192ae5d8c
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/e7bd4f139529e30fb5984d58ee5a942c differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/e999bffefd5fc2e2b14f17fcf0bfda7d b/me-full_25000_5/ocdbt.process_0/d/e999bffefd5fc2e2b14f17fcf0bfda7d
new file mode 100644
index 0000000000000000000000000000000000000000..5cc204ad9e61555f538f75fbfcc6f433cfb4d41d
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/e999bffefd5fc2e2b14f17fcf0bfda7d differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/e9ac5b57affd0d1cc0d2044ec2b54112 b/me-full_25000_5/ocdbt.process_0/d/e9ac5b57affd0d1cc0d2044ec2b54112
new file mode 100644
index 0000000000000000000000000000000000000000..134662a13e513b645a85222cb2167db7edaa3f5a
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/e9ac5b57affd0d1cc0d2044ec2b54112 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/eb0249e75afab7c2d968841f9d955afe b/me-full_25000_5/ocdbt.process_0/d/eb0249e75afab7c2d968841f9d955afe
new file mode 100644
index 0000000000000000000000000000000000000000..f15fc634827edbe14d2bba5b1dbb238780b3bc2c
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/eb0249e75afab7c2d968841f9d955afe differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/eb2b2333e6e4fcfa1f72e16fb0ef5141 b/me-full_25000_5/ocdbt.process_0/d/eb2b2333e6e4fcfa1f72e16fb0ef5141
new file mode 100644
index 0000000000000000000000000000000000000000..fbf68e7b02acc9afa53980aef7b9ce16f8149fe7
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/eb2b2333e6e4fcfa1f72e16fb0ef5141 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/ec960ecbfabb70676403463aa511cc63 b/me-full_25000_5/ocdbt.process_0/d/ec960ecbfabb70676403463aa511cc63
new file mode 100644
index 0000000000000000000000000000000000000000..f1b8cd18403d98c6d8959839c37e96ce5cce2faf
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/ec960ecbfabb70676403463aa511cc63 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/ed8d4ed93318fc43969bceae816a4afb b/me-full_25000_5/ocdbt.process_0/d/ed8d4ed93318fc43969bceae816a4afb
new file mode 100644
index 0000000000000000000000000000000000000000..0873b318e4e0ad1c93b76877289e9d5f4aec5ec0
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/ed8d4ed93318fc43969bceae816a4afb differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/eee84bb0c72c3603f10950bf45b73720 b/me-full_25000_5/ocdbt.process_0/d/eee84bb0c72c3603f10950bf45b73720
new file mode 100644
index 0000000000000000000000000000000000000000..f22d12214e60c53c25f4e1790210efaa40e47b3a
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/eee84bb0c72c3603f10950bf45b73720 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/f35753e062b88fddce59a1aa612487c8 b/me-full_25000_5/ocdbt.process_0/d/f35753e062b88fddce59a1aa612487c8
new file mode 100644
index 0000000000000000000000000000000000000000..fa064ec5d68161a6ede84489892747aee669561b
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/f35753e062b88fddce59a1aa612487c8 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/f418b9c2e64e92f30128df49b2358734 b/me-full_25000_5/ocdbt.process_0/d/f418b9c2e64e92f30128df49b2358734
new file mode 100644
index 0000000000000000000000000000000000000000..4cb5df7989883e391f7275c36ddb76bbd7fbbc41
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/f418b9c2e64e92f30128df49b2358734 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/f61976f1824f129fd80bb3e77d357662 b/me-full_25000_5/ocdbt.process_0/d/f61976f1824f129fd80bb3e77d357662
new file mode 100644
index 0000000000000000000000000000000000000000..9bbc4b532f2f702b07697b8e7475d4677d056a49
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/f61976f1824f129fd80bb3e77d357662 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/f7349119bcd07f48ce447fc7dfeca6f9 b/me-full_25000_5/ocdbt.process_0/d/f7349119bcd07f48ce447fc7dfeca6f9
new file mode 100644
index 0000000000000000000000000000000000000000..57ff48b015c3161e5d892a5dc33cc9ddcfaccbe6
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/f7349119bcd07f48ce447fc7dfeca6f9 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/f738b34b1028398a39183297bbcbaac7 b/me-full_25000_5/ocdbt.process_0/d/f738b34b1028398a39183297bbcbaac7
new file mode 100644
index 0000000000000000000000000000000000000000..ed5fedf95a54fc97cb78deda0d2c51f8dacd5651
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/f738b34b1028398a39183297bbcbaac7 differ
diff --git a/me-full_25000_5/ocdbt.process_0/d/ffb661fc5f4e9ef280c5ba272acab19c b/me-full_25000_5/ocdbt.process_0/d/ffb661fc5f4e9ef280c5ba272acab19c
new file mode 100644
index 0000000000000000000000000000000000000000..8e440e6a0224faa2338d0b72b16bd0237f6a4dfb
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/d/ffb661fc5f4e9ef280c5ba272acab19c differ
diff --git a/me-full_25000_5/ocdbt.process_0/manifest.ocdbt b/me-full_25000_5/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..1a0d25611b969b57f222fb8f4c4e6f684faee74f
Binary files /dev/null and b/me-full_25000_5/ocdbt.process_0/manifest.ocdbt differ